chore: import upstream snapshot with attribution

This commit is contained in:
wehub-resource-sync
2026-07-13 12:47:05 +08:00
commit 4f3b7da785
7394 changed files with 2005594 additions and 0 deletions
@@ -0,0 +1,164 @@
/* ******************************************************************************
*
*
* 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)
//
#include <ops/declarable/helpers/convolutions.h>
#include "cudnnUtils.h"
namespace sd {
namespace ops {
namespace platforms {
//////////////////////////////////////////////////////////////////////////
PLATFORM_IMPL(avgpool2d, ENGINE_CUDA) {
auto input = INPUT_VARIABLE(0);
auto output = OUTPUT_VARIABLE(0);
// 0,1 - kernel Height/Width; 2,3 - stride Height/Width; 4,5 - pad Height/Width; 6,7 - dilation Height/Width; 8 - same
// mode;
const LongType kH = INT_ARG(0);
const LongType kW = INT_ARG(1);
const LongType sH = INT_ARG(2);
const LongType sW = INT_ARG(3);
LongType pH = INT_ARG(4);
LongType pW = INT_ARG(5);
const LongType dH = INT_ARG(6);
const LongType dW = INT_ARG(7);
const auto paddingMode = static_cast<bool>(INT_ARG(8));
const auto extraParam0 = INT_ARG(9);
const int isNCHW = block.getIArguments()->size() > 10 ? !INT_ARG(10) : 1; // INT_ARG(10): 0-NCHW, 1-NHWC
REQUIRE_TRUE(input->rankOf() == 4, 0, "AVGPOOL2D CUDNN op: input should have rank of 4, but got %i instead",
input->rankOf());
REQUIRE_TRUE(dH != 0 && dW != 0, 0, "AVGPOOL2D CUDNN op: dilation must not be zero, but got instead {%i, %i}", dH,
dW);
LongType oH = 0;
LongType oW = 0;
const LongType iH = static_cast<LongType>(isNCHW ? input->sizeAt(2) : input->sizeAt(1));
const LongType iW = static_cast<LongType>(isNCHW ? input->sizeAt(3) : input->sizeAt(2));
ConvolutionUtils::calcOutSizePool2D(oH, oW, kH, kW, sH, sW, pH, pW, dH, dW, iH, iW, paddingMode);
if (paddingMode) ConvolutionUtils::calcPadding2D(pH, pW, oH, oW, iH, iW, kH, kW, sH, sW, dH, dW);
const cudnnPoolingMode_t mode =
(extraParam0 == 0) ? CUDNN_POOLING_AVERAGE_COUNT_EXCLUDE_PADDING : CUDNN_POOLING_AVERAGE_COUNT_INCLUDE_PADDING;
pooling2dCUDNN(block.launchContext(), input, output, kH, kW, sH, sW, pH, pW, dH, dW, isNCHW, mode);
return Status::OK;
}
//////////////////////////////////////////////////////////////////////////
PLATFORM_CHECK(avgpool2d, ENGINE_CUDA) {
auto input = INPUT_VARIABLE(0);
auto output = OUTPUT_VARIABLE(0);
Requirements req("CUDNN AVGPOOL2d OP");
req.expectEq(makeInfoVariable(input->dataType(), TYPE_MSG_INPUT),
makeInfoVariable(output->dataType(), TYPE_MSG_OUTPUT)) &&
req.expectIn(makeInfoVariable(input->dataType(), TYPE_MSG_INPUT),
{INT32, HALF, FLOAT32, DOUBLE});
req.logTheSuccess();
return req;
}
//////////////////////////////////////////////////////////////////////////
PLATFORM_IMPL(avgpool2d_bp, ENGINE_CUDA) {
auto input = INPUT_VARIABLE(0); // [bS, iH, iW, iC] (NHWC) or [bS, iC, iH, iW] (NCHW)
auto gradO = INPUT_VARIABLE(1); // [bS, oH, oW, oC] (NHWC) or [bS, oC, oH, oW] (NCHW), epsilon_next
auto gradI = OUTPUT_VARIABLE(0); // [bS, iH, iW, iC] (NHWC) or [bS, iC, iH, iW] (NCHW), epsilon
const LongType kH = INT_ARG(0); // filter(kernel) height
const LongType kW = INT_ARG(1); // filter(kernel) width
const LongType sH = INT_ARG(2); // strides height
const LongType sW = INT_ARG(3); // strides width
LongType pH = INT_ARG(4); // paddings height
LongType pW = INT_ARG(5); // paddings width
const LongType dH = INT_ARG(6); // dilations height
const LongType dW = INT_ARG(7); // dilations width
const auto paddingMode = INT_ARG(8); // 0-VALID, 1-SAME
const auto extraParam0 = INT_ARG(9);
const auto isNCHW = block.getIArguments()->size() > 10 ? !INT_ARG(10) : 1; // INT_ARG(10): 0-NCHW, 1-NHWC
REQUIRE_TRUE(input->rankOf() == 4, 0, "AVGPOOL2D_BP CUDNN op: input should have rank of 4, but got %i instead",
input->rankOf());
REQUIRE_TRUE(dH != 0 && dW != 0, 0, "AVGPOOL2D_BP CUDNN op: dilation must not be zero, but got instead {%i, %i}", dH,
dW);
LongType bS, iC, iH, iW, oC, oH,
oW; // batch size, input channels, input height/width, output channels, output height/width;
LongType indIOioC, indIiH, indWoC, indWiC, indWkH, indOoH; // corresponding indexes
ConvolutionUtils::getSizesAndIndexesConv2d(isNCHW, 0, *input, *gradO, bS, iC, iH, iW, oC, oH, oW, indIOioC, indIiH,
indWiC, indWoC, indWkH, indOoH);
std::vector<LongType> expectedGradOShape =
ShapeUtils::composeShapeUsingDimsAndIdx({bS, iC, oH, oW, 0, indIOioC, indIiH, indIiH + 1});
std::vector<LongType> expectedGradIShape =
ShapeUtils::composeShapeUsingDimsAndIdx({bS, iC, iH, iW, 0, indIOioC, indIiH, indIiH + 1});
REQUIRE_TRUE(gradO->isSameShape(expectedGradOShape), 0,
"AVGPOOL2D_BP CUDNN op: wrong shape of output's gradients array (next epsilon), expected is %s, but got "
"%s instead !",
ShapeUtils::shapeAsString(expectedGradOShape).c_str(), ShapeUtils::shapeAsString(gradO).c_str());
REQUIRE_TRUE(
gradI->isSameShape(expectedGradIShape), 0,
"AVGPOOL2D_BP CUDNN op: wrong shape of input's gradients array (epsilon), expected is %s, but got %s instead !",
ShapeUtils::shapeAsString(expectedGradIShape).c_str(), ShapeUtils::shapeAsString(gradI).c_str());
if (paddingMode) // SAME
ConvolutionUtils::calcPadding2D(pH, pW, oH, oW, iH, iW, kH, kW, sH, sW, dH, dW);
const cudnnPoolingMode_t mode =
(extraParam0 == 0) ? CUDNN_POOLING_AVERAGE_COUNT_EXCLUDE_PADDING : CUDNN_POOLING_AVERAGE_COUNT_INCLUDE_PADDING;
pooling2dBpCUDNN(block.launchContext(), input, gradO, gradI, kH, kW, sH, sW, pH, pW, dH, dW, isNCHW, mode);
return Status::OK;
}
PLATFORM_CHECK(avgpool2d_bp, ENGINE_CUDA) {
auto input = INPUT_VARIABLE(0); // [bS, iH, iW, iC] (NHWC) or [bS, iC, iH, iW] (NCHW)
auto gradO = INPUT_VARIABLE(1); // [bS, oH, oW, oC] (NHWC) or [bS, oC, oH, oW] (NCHW), epsilon_next
auto gradI = OUTPUT_VARIABLE(0); // [bS, iH, iW, iC] (NHWC) or [bS, iC, iH, iW] (NCHW), epsilon
Requirements req("CUDNN AVGPOOL2d_BP OP");
req.expectEq(makeInfoVariable(input->dataType(), TYPE_MSG_INPUT0),
makeInfoVariable(gradO->dataType(), TYPE_MSG_INPUT1)) &&
req.expectEq(makeInfoVariable(input->dataType(), TYPE_MSG_INPUT),
makeInfoVariable(gradI->dataType(), TYPE_MSG_OUTPUT)) &&
req.expectIn(makeInfoVariable(input->dataType(), TYPE_MSG_INPUT),
{INT32, HALF, FLOAT32, DOUBLE}) &&
req.expect(
makeShapeInfoVariable(input, SHAPE_MSG_INPUT0), makeShapeInfoVariable(gradI, SHAPE_MSG_OUTPUT),
[](const decltype(input)& l, const decltype(gradI)& r) {
return shape::haveSameShapeAndStrides(l->shapeInfo(), r->shapeInfo());
},
EXPECTED_EQ_MSG);
req.logTheSuccess();
return req;
}
} // namespace platforms
} // namespace ops
} // namespace sd
@@ -0,0 +1,176 @@
/* ******************************************************************************
*
*
* 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)
//
#include <ops/declarable/helpers/convolutions.h>
#include "cudnnUtils.h"
namespace sd {
namespace ops {
namespace platforms {
//////////////////////////////////////////////////////////////////////////
PLATFORM_IMPL(avgpool3dnew, ENGINE_CUDA) {
auto input = INPUT_VARIABLE(0); // [bS, iD, iH, iW, iC] (NDHWC) or [bS, iC, iD, iH, iW] (NCDHW)
auto output = OUTPUT_VARIABLE(0); // [bS, oD, oH, oW, iC] (NDHWC) or [bS, iC, oD, oH, oW] (NCDHW)
LongType kD = INT_ARG(0); // filter(kernel) depth
LongType kH = INT_ARG(1); // filter(kernel) height
LongType kW = INT_ARG(2); // filter(kernel) width
LongType sD = INT_ARG(3); // strides depth
LongType sH = INT_ARG(4); // strides height
LongType sW = INT_ARG(5); // strides width
LongType pD = INT_ARG(6); // paddings depth
LongType pH = INT_ARG(7); // paddings height
LongType pW = INT_ARG(8); // paddings width
LongType dD = INT_ARG(9); // dilations depth
LongType dH = INT_ARG(10); // dilations height
LongType dW = INT_ARG(11); // dilations width
int paddingMode = INT_ARG(12); // 1-SAME, 0-VALID
int extraParam0 = INT_ARG(13);
int isNCDHW = block.getIArguments()->size() > 14 ? !INT_ARG(14) : 1; // 0-NCDHW, 1-NDHWC
REQUIRE_TRUE(input->rankOf() == 5, 0,
"AVGPOOL3DNEW CUDNN OP: rank of input array must be equal to 5, but got %i instead !", input->rankOf());
REQUIRE_TRUE(dD != 0 && dH != 0 && dW != 0, 0,
"AVGPOOL3DNEW CUDNN OP: dilation must not be zero, but got instead {%i, %i, %i}", dD, dH, dW);
LongType bS, iC, iD, iH, iW, oC, oD, oH,
oW; // batch size, input channels, input depth/height/width, output channels, output depth/height/width;
LongType indIOioC, indIOioD, indWoC, indWiC, indWkD; // corresponding indexes
ConvolutionUtils::getSizesAndIndexesConv3d(isNCDHW, 0, *input, *output, bS, iC, iD, iH, iW, oC, oD, oH, oW, indIOioC,
indIOioD, indWiC, indWoC, indWkD);
std::vector<LongType> expectedOutputShape =
ShapeUtils::composeShapeUsingDimsAndIdx({bS, iC, oD, oH, oW, 0, indIOioC, indIOioD, indIOioD + 1, indIOioD + 2});
REQUIRE_TRUE(output->isSameShape(expectedOutputShape), 0,
"AVGPOOL3DNEW CUDNN OP: wrong shape of output array, expected is %s, but got %s instead !",
ShapeUtils::shapeAsString(expectedOutputShape).c_str(), ShapeUtils::shapeAsString(output).c_str());
if (paddingMode) // SAME
ConvolutionUtils::calcPadding3D(pD, pH, pW, oD, oH, oW, iD, iH, iW, kD, kH, kW, sD, sH, sW, dD, dH, dW);
const cudnnPoolingMode_t mode =
(extraParam0 == 0) ? CUDNN_POOLING_AVERAGE_COUNT_EXCLUDE_PADDING : CUDNN_POOLING_AVERAGE_COUNT_INCLUDE_PADDING;
pooling3dCUDNN(block.launchContext(), input, output, kD, kH, kW, sD, sH, sW, pD, pH, pW, dD, dH, dW, isNCDHW, mode);
return Status::OK;
}
//////////////////////////////////////////////////////////////////////////
PLATFORM_CHECK(avgpool3dnew, ENGINE_CUDA) {
auto input = INPUT_VARIABLE(0);
auto output = OUTPUT_VARIABLE(0);
Requirements req("CUDNN AVGPOOL3d OP");
req.expectEq(makeInfoVariable(input->dataType(), TYPE_MSG_INPUT),
makeInfoVariable(output->dataType(), TYPE_MSG_OUTPUT)) &&
req.expectIn(makeInfoVariable(input->dataType(), TYPE_MSG_INPUT),
{INT32, HALF, FLOAT32, DOUBLE});
req.logTheSuccess();
return req;
}
//////////////////////////////////////////////////////////////////////////
PLATFORM_IMPL(avgpool3dnew_bp, ENGINE_CUDA) {
auto input = INPUT_VARIABLE(0); // [bS, iD, iH, iW, iC] (NDHWC) or [bS, iC, iD, iH, iW] (NCDHW)
auto gradO = INPUT_VARIABLE(1); // [bS, oD, oH, oW, oC] (NDHWC) or [bS, oC, oD, oH, oW] (NCDHW), epsilon_next
auto gradI = OUTPUT_VARIABLE(0); // [bS, iD, iH, iW, iC] (NDHWC) or [bS, iC, iD, iH, iW] (NCDHW), epsilon
const LongType kD = INT_ARG(0); // filter(kernel) depth
const LongType kH = INT_ARG(1); // filter(kernel) height
const LongType kW = INT_ARG(2); // filter(kernel) width
const LongType sD = INT_ARG(3); // strides depth
const LongType sH = INT_ARG(4); // strides height
const LongType sW = INT_ARG(5); // strides width
LongType pD = INT_ARG(6); // paddings depth
LongType pH = INT_ARG(7); // paddings height
LongType pW = INT_ARG(8); // paddings width
const LongType dD = INT_ARG(9); // dilations depth
const LongType dH = INT_ARG(10); // dilations height
const LongType dW = INT_ARG(11); // dilations width
const int isSameMode = INT_ARG(12); // 1-SAME, 0-VALID
const int extraParam0 = INT_ARG(13); // define what divisor to use while averaging
const int isNCDHW = block.getIArguments()->size() > 14 ? !INT_ARG(14) : 1; // 0-NCDHW, 1-NDHWC
REQUIRE_TRUE(input->rankOf() == 5, 0, "AVGPOOL3DNEW_BP CUDNN OP: input should have rank of 5, but got %i instead",
input->rankOf());
REQUIRE_TRUE(dD != 0 && dH != 0 && dW != 0, 0,
"AVGPOOL3DNEW_BP CUDNN OP: dilation must not be zero, but got instead {%i, %i, %i}", dD, dH, dW);
LongType bS, iC, iD, iH, iW, oC, oD, oH,
oW; // batch size, input channels, input depth/height/width, output channels, output depth/height/width;
LongType indIOioC, indIOioD, indWoC, indWiC, indWkD; // corresponding indexes
ConvolutionUtils::getSizesAndIndexesConv3d(isNCDHW, 0, *input, *gradO, bS, iC, iD, iH, iW, oC, oD, oH, oW, indIOioC,
indIOioD, indWiC, indWoC, indWkD);
std::vector<LongType> expectedGradOShape =
ShapeUtils::composeShapeUsingDimsAndIdx({bS, iC, oD, oH, oW, 0, indIOioC, indIOioD, indIOioD + 1, indIOioD + 2});
std::vector<LongType> expectedGradIShape =
ShapeUtils::composeShapeUsingDimsAndIdx({bS, iC, iD, iH, iW, 0, indIOioC, indIOioD, indIOioD + 1, indIOioD + 2});
REQUIRE_TRUE(gradO->isSameShape(expectedGradOShape), 0,
"AVGPOOL3DNEW_BP CUDNN: wrong shape of output's gradients array (next epsilon), expected is %s, but got "
"%s instead !",
ShapeUtils::shapeAsString(expectedGradOShape).c_str(), ShapeUtils::shapeAsString(gradO).c_str());
REQUIRE_TRUE(
gradI->isSameShape(expectedGradIShape), 0,
"AVGPOOL3DNEW_BP CUDNN: wrong shape of input's gradients array (epsilon), expected is %s, but got %s instead !",
ShapeUtils::shapeAsString(expectedGradIShape).c_str(), ShapeUtils::shapeAsString(gradI).c_str());
if (isSameMode) // SAME
ConvolutionUtils::calcPadding3D(pD, pH, pW, oD, oH, oW, iD, iH, iW, kD, kH, kW, sD, sH, sW, dD, dH, dW);
const cudnnPoolingMode_t mode =
(extraParam0 == 0) ? CUDNN_POOLING_AVERAGE_COUNT_EXCLUDE_PADDING : CUDNN_POOLING_AVERAGE_COUNT_INCLUDE_PADDING;
pooling3dBpCUDNN(block.launchContext(), input, gradO, gradI, kD, kH, kW, sD, sH, sW, pD, pH, pW, dD, dH, dW, isNCDHW,
mode);
return Status::OK;
}
PLATFORM_CHECK(avgpool3dnew_bp, ENGINE_CUDA) {
auto input = INPUT_VARIABLE(0); // [bS, iD, iH, iW, iC] (NDHWC) or [bS, iC, iD, iH, iW] (NCDHW)
auto gradO = INPUT_VARIABLE(1); // [bS, oD, oH, oW, oC] (NDHWC) or [bS, oC, oD, oH, oW] (NCDHW), epsilon_next
auto gradI = OUTPUT_VARIABLE(0); // [bS, iD, iH, iW, iC] (NDHWC) or [bS, iC, iD, iH, iW] (NCDHW), epsilon
Requirements req("CUDNN AVGPOOL3d_BP OP");
req.expectEq(makeInfoVariable(input->dataType(), TYPE_MSG_INPUT0),
makeInfoVariable(gradO->dataType(), TYPE_MSG_INPUT1)) &&
req.expectEq(makeInfoVariable(input->dataType(), TYPE_MSG_INPUT),
makeInfoVariable(gradI->dataType(), TYPE_MSG_OUTPUT)) &&
req.expectIn(makeInfoVariable(input->dataType(), TYPE_MSG_INPUT),
{INT32, HALF, FLOAT32, DOUBLE}) &&
req.expect(
makeShapeInfoVariable(input, SHAPE_MSG_INPUT0), makeShapeInfoVariable(gradI, SHAPE_MSG_OUTPUT),
[](const decltype(input)& l, const decltype(gradI)& r) {
return shape::haveSameShapeAndStrides(l->shapeInfo(), r->shapeInfo());
},
EXPECTED_EQ_MSG);
req.logTheSuccess();
return req;
}
} // namespace platforms
} // namespace ops
} // namespace sd
@@ -0,0 +1,557 @@
/* ******************************************************************************
*
*
* 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)
//
#include <ops/declarable/helpers/convolutions.h>
#include "cudnnUtils.h"
namespace sd {
namespace ops {
namespace platforms {
//////////////////////////////////////////////////////////////////////////
static void batchnormCUDNN(const LaunchContext* context, NDArray* input, NDArray* mean,
NDArray* variance, NDArray* gamma, NDArray* beta, NDArray* output,
const double epsilon, const bool isSpatialMode) {
// input, output -> 4D:nchw, 5D:ncdhw
// mean, variance, gamma, beta -> 1xCx1x1 for 4D and 1xCx1x1x1 for 5D for BATCHNORM_MODE_SPATIAL mode
// -> 1xCxHxW for 4D and 1xCxDxHxW for 5D for BATCHNORM_MODE_PER_ACTIVATION mode
const cudnnDataType_t dataType = cudnnDataType(input->dataType());
const LongType xRank = input->rankOf();
auto handle = reinterpret_cast<cudnnHandle_t*>(context->getCuDnnHandle());
CHECK_CUDNN_FAILURE(cudnnSetStream(*handle, *context->getCudaStream()));
const std::vector<int> xShape = input->getShapeAsVectorInt(); // input and output have same shapes
std::vector<int> paramsShape, paramsStrides; // mean, variance, gamma and beta have same shapes
if (isSpatialMode) { // 1xCx1x1
const int iC = static_cast<int>(mean->lengthOf());
const int stride0 = static_cast<int>(mean->strideAt(0));
paramsShape = xRank == 4 ? std::vector<int>({1, iC, 1, 1}) : std::vector<int>({1, iC, 1, 1, 1});
paramsStrides = xRank == 4 ? std::vector<int>({iC * stride0, stride0, 1, 1})
: std::vector<int>({iC * stride0, stride0, 1, 1, 1});
} else {
paramsShape = std::vector<int>(mean->getShapeAsVector().begin(), mean->getShapeAsVector().end());
paramsStrides = xRank == 4
? std::vector<int>({static_cast<int>(mean->strideAt(0)), static_cast<int>(mean->strideAt(1)), static_cast<int>(mean->strideAt(2)),
static_cast<int>(mean->strideAt(3))})
: std::vector<int>({static_cast<int>(mean->strideAt(0)), static_cast<int>(mean->strideAt(1)), static_cast<int>(mean->strideAt(2)),
static_cast<int>(mean->strideAt(3)), static_cast<int>(mean->strideAt(4))});
}
std::vector<int> xStrides = {static_cast<int>(input->strideAt(0)), static_cast<int>(input->strideAt(1)), static_cast<int>(input->strideAt(2)),
static_cast<int>(input->strideAt(3))};
std::vector<int> zStrides = {static_cast<int>(output->strideAt(0)), static_cast<int>(output->strideAt(1)), static_cast<int>(output->strideAt(2)),
static_cast<int>(output->strideAt(3))};
if (xRank > 4) { // 5D
xStrides.push_back((LongType)input->strideAt(4));
zStrides.push_back((LongType)output->strideAt(4));
}
cudnnTensorFormat_t format = CUDNN_TENSOR_NCHW;
// input descriptor
x.set(dataType, xRank, xShape.data(), xStrides.data());
// output descriptor
CudnnTensor z;
z.set(dataType, xRank, xShape.data(), zStrides.data());
// mean, variance, gamma and beta descriptor, the same descriptor for all of them
CudnnTensor params;
params.set(dataType, xRank, paramsShape.data(), paramsStrides.data());
// provide scaling parameters
const float alpha32(1), beta32(0);
const double alpha64(1), beta64(0);
const void* ptrAlpha =
output->sizeOfT() <= 4 ? reinterpret_cast<const void*>(&alpha32) : reinterpret_cast<const void*>(&alpha64);
const void* ptrBeta =
output->sizeOfT() <= 4 ? reinterpret_cast<const void*>(&beta32) : reinterpret_cast<const void*>(&beta64);
NDArray::prepareSpecialUse({output}, {input, mean, variance, gamma, beta});
// calculations
CHECK_CUDNN_FAILURE_MSG(
STRINGIZE(cudnnBatchNormalizationForwardInference),
cudnnBatchNormalizationForwardInference(
*handle, isSpatialMode ? CUDNN_BATCHNORM_SPATIAL : CUDNN_BATCHNORM_PER_ACTIVATION, ptrAlpha, ptrBeta, x,
input->specialBuffer(), z, output->specialBuffer(), params, gamma->specialBuffer(), beta->specialBuffer(),
mean->specialBuffer(), variance->specialBuffer(), epsilon));
auto cudaErr = cudaStreamSynchronize(*context->getCudaStream());
if (cudaErr != 0) throw cuda_exception::build("batchnormCUDNN: cudaStreamSynchronize failed !", cudaErr);
NDArray::registerSpecialUse({output}, {input, mean, variance, gamma, beta});
}
//////////////////////////////////////////////////////////////////////////
static void batchnormBpCUDNN(const LaunchContext* context, NDArray* input, NDArray* mean,
NDArray* variance, NDArray* gamma, NDArray* gradO, NDArray* gradI,
NDArray* gradG, NDArray* gradB, const double epsilon, const bool isSpatialMode) {
// input, gradO, gradI -> 4D:nchw, 5D:ncdhw
// mean, variance, gamma, beta, gradM, gradV, gradG, gradB -> 1xCx1x1 for 4D and 1xCx1x1x1 for 5D for
// BATCHNORM_MODE_SPATIAL mode
// -> 1xCxHxW for 4D and 1xCxDxHxW for 5D for
// BATCHNORM_MODE_PER_ACTIVATION mode
const cudnnDataType_t dataType = cudnnDataType(input->dataType());
const int xRank = input->rankOf();
auto handle = reinterpret_cast<cudnnHandle_t*>(context->getCuDnnHandle());
cudnnStatus_t err = cudnnSetStream(*handle, *context->getCudaStream());
const std::vector<int> xShape = input->getShapeAsVectorInt(); // input and output have same shapes
std::vector<int> paramsShape, paramsStrides; // mean, variance, gamma and beta have same shapes
if (isSpatialMode) { // 1xCx1x1
const int iC = static_cast<int>(mean->lengthOf());
const int stride0 = static_cast<int>(mean->strideAt(0));
paramsShape = xRank == 4 ? std::vector<int>({1, iC, 1, 1}) : std::vector<int>({1, iC, 1, 1, 1});
paramsStrides = xRank == 4 ? std::vector<int>({iC * stride0, stride0, 1, 1})
: std::vector<int>({iC * stride0, stride0, 1, 1, 1});
} else {
paramsShape = std::vector<int>(mean->getShapeAsVector().begin(), mean->getShapeAsVector().end());
paramsStrides = xRank == 4
? std::vector<int>({static_cast<int>(mean->strideAt(0)), static_cast<int>(mean->strideAt(1)), static_cast<int>(mean->strideAt(2)),
static_cast<int>(mean->strideAt(3))})
: std::vector<int>({static_cast<int>(mean->strideAt(0)), static_cast<int>(mean->strideAt(1)), static_cast<int>(mean->strideAt(2)),
static_cast<int>(mean->strideAt(3)), static_cast<int>(mean->strideAt(4))});
}
std::vector<int> xStrides = {static_cast<int>(input->strideAt(0)), static_cast<int>(input->strideAt(1)), static_cast<int>(input->strideAt(2)),
static_cast<int>(input->strideAt(3))};
std::vector<int> dxStrides = {static_cast<int>(gradI->strideAt(0)), static_cast<int>(gradI->strideAt(1)), static_cast<int>(gradI->strideAt(2)),
static_cast<int>(gradI->strideAt(3))};
std::vector<int> dzStrides = {static_cast<int>(gradO->strideAt(0)), static_cast<int>(gradO->strideAt(1)), static_cast<int>(gradO->strideAt(2)),
static_cast<int>(gradO->strideAt(3))};
if (xRank > 4) { // 5D
xStrides.push_back(static_cast<int>(input->strideAt(4)));
dxStrides.push_back(static_cast<int>(gradI->strideAt(4)));
dzStrides.push_back(static_cast<int>(gradO->strideAt(4)));
}
cudnnTensorFormat_t format = CUDNN_TENSOR_NCHW;
// input descriptor
CudnnTensor x;
x.set(dataType, xRank, xShape.data(), xStrides.data());
// gradO descriptor
CudnnTensor dz;
dz.set(dataType, xRank, xShape.data(), dzStrides.data());
// gradI descriptor
CudnnTensor dx;
dx.set(dataType, xRank, xShape.data(), dxStrides.data());
// mean, variance, gamma, gradG and gradB descriptor, the same descriptor for all of them
CudnnTensor params;
params.set(dataType, xRank, paramsShape.data(), paramsStrides.data());
// provide scaling parameters
const float alpha32(1), beta32(0);
double alpha64(1), beta64(0);
const void* ptrAlpha =
input->sizeOfT() <= 4 ? reinterpret_cast<const void*>(&alpha32) : reinterpret_cast<const void*>(&alpha64);
const void* ptrBeta =
input->sizeOfT() <= 4 ? reinterpret_cast<const void*>(&beta32) : reinterpret_cast<const void*>(&beta64);
NDArray::prepareSpecialUse({gradI, gradG, gradB}, {input, mean, variance, gamma, gradO});
// calculations
// TODO: we can use cache here
CHECK_CUDNN_FAILURE_MSG(
STRINGIZE(cudnnBatchNormalizationBackward),
cudnnBatchNormalizationBackward(*handle, isSpatialMode ? CUDNN_BATCHNORM_SPATIAL : CUDNN_BATCHNORM_PER_ACTIVATION,
ptrAlpha, ptrBeta, ptrAlpha, ptrBeta, x, input->specialBuffer(), dz,
gradO->specialBuffer(), dx, gradI->specialBuffer(), params,
gamma->specialBuffer(), gradG->specialBuffer(), gradB->specialBuffer(), epsilon,
nullptr /*mean->specialBuffer()*/, nullptr /*variance->specialBuffer()*/));
auto cudaErr = cudaStreamSynchronize(*context->getCudaStream());
if (cudaErr != 0) throw cuda_exception::build("batchnormBpCUDNN: cudaStreamSynchronize failed !", cudaErr);
NDArray::registerSpecialUse({gradI, gradG, gradB}, {input, mean, variance, gamma, gradO});
}
//////////////////////////////////////////////////////////////////////////
PLATFORM_IMPL(batchnorm, ENGINE_CUDA) {
auto input = INPUT_VARIABLE(0);
auto mean = INPUT_VARIABLE(1);
auto variance = INPUT_VARIABLE(2);
NDArray* gamma = nullptr;
NDArray* beta = nullptr;
auto output = OUTPUT_VARIABLE(0);
const bool applyScale = (bool)INT_ARG(0);
const bool applyOffset = (bool)INT_ARG(1);
const double epsilon = T_ARG(0);
if (applyScale) gamma = INPUT_VARIABLE(3);
if (applyOffset) beta = INPUT_VARIABLE(3 + (int)applyScale);
const int numOfIntArgs = block.getIArguments()->size();
const int inRank = input->rankOf();
// get axes args to normalize input array over
std::vector<int> axes;
if (numOfIntArgs > 2)
for (int i = 2; i < numOfIntArgs; ++i) axes.push_back(INT_ARG(i));
else
axes.push_back(inRank - 1); // default dimension to reduce along is last dimension
const int numOfAxes = axes.size();
REQUIRE_TRUE(numOfAxes <= inRank, 0,
"BATCHNORM CUDNN op: too big number of input axes to normalize over, expected number should be less or "
"equal to rank of input array, but got %i and %i correspondingly !",
numOfAxes, inRank);
// evaluate expected shape for mean, variance and gamma. These 3 arrays should have identical shapes
// for example if input shape is {2,3,4,5,6} and axes = {1,3}, then expected shape would be {1,3,1,5,1}, and if axes =
// {3}, then expected shape would be {5}
std::vector<LongType> expShape;
if (numOfAxes == 1)
expShape.push_back(input->sizeAt(axes[0]));
else { // get, for example, something like {1, inputDim1, 1, inputDim3, 1} if axes = {1, 3}
expShape = std::vector<LongType>(inRank, 1);
for (LongType i = 0; i < numOfAxes; ++i) expShape[axes[i]] = input->sizeAt(axes[i]);
}
REQUIRE_TRUE(mean->isSameShape(expShape), 0,
"BATCHNORM CUDNN op: wrong shape of mean array, expected is %s, but got %s instead !",
ShapeUtils::shapeAsString(expShape).c_str(), ShapeUtils::shapeAsString(mean).c_str());
REQUIRE_TRUE(variance->isSameShape(expShape), 0,
"BATCHNORM CUDNN op: wrong shape of variance array, expected is %s, but got %s instead !",
ShapeUtils::shapeAsString(expShape).c_str(), ShapeUtils::shapeAsString(variance).c_str());
if (gamma)
REQUIRE_TRUE(gamma->isSameShape(expShape), 0,
"BATCHNORM CUDNN op: wrong shape of gamma array, expected is %s, but got %s instead !",
ShapeUtils::shapeAsString(expShape).c_str(), ShapeUtils::shapeAsString(gamma).c_str());
if (beta)
REQUIRE_TRUE(beta->isSameShape(expShape), 0,
"BATCHNORM CUDNN op: wrong shape of beta array, expected is %s, but got %s instead !",
ShapeUtils::shapeAsString(expShape).c_str(), ShapeUtils::shapeAsString(beta).c_str());
// types of all input arrays should be the same
for (int i = 1; i < block.width(); ++i)
REQUIRE_TRUE(INPUT_VARIABLE(0)->dataType() == INPUT_VARIABLE(i)->dataType(), 0,
"BATCHNORM CUDNN op: types of all input arrays should be the same !");
// cudnn supports NCHW format only
const bool needPermut = axes.size() == 1 && mean->lengthOf() == input->sizeAt(-1);
std::unique_ptr<NDArray> tmpGamma = {}, tmpBeta = {}, tmpInput = {}, tmpOutput = {};
if (needPermut) { // if NHWC
std::vector<LongType> perm =
inRank == 4 ? std::vector<LongType>({0, 3, 1, 2}) : std::vector<LongType>({0, 4, 1, 2, 3}); // NHWC -> NCHW
tmpInput.reset(new NDArray(input->permute(perm)));
tmpOutput.reset(new NDArray(output->permute(perm)));
input = tmpInput.get();
output = tmpOutput.get();
}
// cudnn requires gamma and beta to be non-nullptr
if (!applyScale) {
tmpGamma.reset(new NDArray(mean));
gamma = tmpGamma.get();
*gamma = 1;
}
if (!applyOffset) {
tmpBeta.reset(new NDArray(mean));
beta = tmpBeta.get();
*beta = 0;
}
// calculations
batchnormCUDNN(block.launchContext(), input, mean, variance, gamma, beta, output, epsilon, axes.size() == 1);
return Status::OK;
}
//////////////////////////////////////////////////////////////////////////
PLATFORM_CHECK(batchnorm, ENGINE_CUDA) {
const bool applyScale = (bool)INT_ARG(0);
const bool applyOffset = (bool)INT_ARG(1);
NDArray* input = INPUT_VARIABLE(0);
NDArray* mean = INPUT_VARIABLE(1);
NDArray* variance = INPUT_VARIABLE(2);
NDArray* gamma = applyScale ? INPUT_VARIABLE(3) : nullptr;
NDArray* beta = applyOffset ? INPUT_VARIABLE(3 + (int)applyScale) : nullptr;
const int numOfIntArgs = block.getIArguments()->size();
const int xRank = input->rankOf();
// *********************************** //
// get axes args to normalize input array over
std::vector<int> axes;
if (numOfIntArgs > 2)
for (int i = 2; i < numOfIntArgs; ++i) axes.push_back(INT_ARG(i));
else
axes.push_back(xRank - 1); // default dimension to reduce along is last dimension
Requirements req("CUDNN BATCHNORM OP");
req.expectIn(makeInfoVariable(xRank, RANK_MSG_INPUT0), {4, 5}) &&
req.expectIn(makeInfoVariable(input->dataType(), TYPE_MSG_INPUT0),
{HALF, FLOAT32, DOUBLE}) &&
req.expectIn(makeInfoVariable(axes.size(), "axes.size()"), {1, 3, 4}) &&
req.expect(
makeShapeInfoVariable(mean, SHAPE_MSG_INPUT1), makeShapeInfoVariable(variance, SHAPE_MSG_INPUT2),
[](const decltype(mean)& l, const decltype(variance)& r) {
return shape::haveSameShapeAndStrides(l->shapeInfo(), r->shapeInfo());
},
EXPECTED_EQ_MSG);
if (gamma) {
req.expect(
makeShapeInfoVariable(gamma, SHAPE_MSG_INPUT_ "#gamma"), makeShapeInfoVariable(mean, SHAPE_MSG_INPUT1),
[](const decltype(gamma)& l, const decltype(mean)& r) {
return shape::haveSameShapeAndStrides(l->shapeInfo(), r->shapeInfo());
},
EXPECTED_EQ_MSG);
}
if (beta) {
req.expect(
makeShapeInfoVariable(beta, SHAPE_MSG_INPUT_ "#beta"), makeShapeInfoVariable(mean, SHAPE_MSG_INPUT1),
[](const decltype(beta)& l, const decltype(mean)& r) {
return shape::haveSameShapeAndStrides(l->shapeInfo(), r->shapeInfo());
},
EXPECTED_EQ_MSG);
}
if (axes.size() == 1) {
req.expectIn(makeInfoVariable(mean->lengthOf(), LENGTH_MSG_INPUT1), {-1, 1});
} else {
auto inputShapeModif = input->getShapeAsVector(); // [dim0,dim1,dim2,dim3] 4D or [dim0,dim1,dim2,dim3,dim4]
inputShapeModif[0] = 1;
// mean [1,dim1,dim2,dim3] 4D or [1,dim1,dim2,dim3,dim4]
req.expect(
makeShapeInfoVariable(mean, SHAPE_MSG_INPUT1),
makeShapeInfoVariable(inputShapeModif, SHAPE_MSG_INPUT_ "#expect"),
[](const decltype(mean)& l, const decltype(inputShapeModif)& r) { return l->isSameShape(r); }, EXPECTED_EQ_MSG);
}
req.logTheSuccess();
return req;
}
//////////////////////////////////////////////////////////////////////////
PLATFORM_IMPL(batchnorm_bp, ENGINE_CUDA) {
NDArray* input = INPUT_VARIABLE(0);
NDArray* mean = INPUT_VARIABLE(1);
NDArray* variance = INPUT_VARIABLE(2);
NDArray* gamma = nullptr;
NDArray* beta = nullptr;
NDArray* gradO = INPUT_VARIABLE(block.width() - 1); // next epsilon
NDArray* gradI = OUTPUT_VARIABLE(0);
NDArray* gradM = OUTPUT_VARIABLE(1);
NDArray* gradV = OUTPUT_VARIABLE(2);
NDArray* gradG = nullptr;
NDArray* gradB = nullptr;
const bool applyScale = (bool)INT_ARG(0);
const bool applyOffset = (bool)INT_ARG(1);
const float epsilon = T_ARG(0);
if (applyScale) {
gamma = INPUT_VARIABLE(3);
gradG = OUTPUT_VARIABLE(3);
}
if (applyOffset) {
beta = INPUT_VARIABLE(3 + (int)applyScale);
gradB = OUTPUT_VARIABLE(3 + (int)applyScale);
}
const int numOfIntArgs = block.getIArguments()->size();
const int inRank = input->rankOf();
// get axes args to normalize input array over
std::vector<int> axes;
if (numOfIntArgs > 2)
for (int i = 2; i < numOfIntArgs; ++i) axes.push_back(INT_ARG(i));
else
axes.push_back(inRank - 1); // default dimension to reduce along is last dimension
const int numOfAxes = axes.size();
REQUIRE_TRUE(numOfAxes <= inRank, 0,
"BATCHNORM_BP CUDNN op: too big number of input axes to normalize over, expected number should be less "
"or equal to rank of input array, but got %i and %i correspondingly !",
numOfAxes, inRank);
// evaluate expected shape for mean, variance and gamma. These 3 arrays should have identical shapes
// for example if input shape is {2,3,4,5,6} and axes = {1,3}, then expected shape would be {1,3,1,5,1}, and if axes =
// {3}, then expected shape would be {5}
std::vector<LongType> expShape;
if (numOfAxes == 1)
expShape.push_back(input->sizeAt(axes[0]));
else { // get, for example, something like {1, inputDim1, 1, inputDim3, 1} if axes = {1, 3}
expShape = std::vector<LongType>(inRank, 1);
for (LongType i = 0; i < numOfAxes; ++i) expShape[axes[i]] = input->sizeAt(axes[i]);
}
REQUIRE_TRUE(mean->isSameShape(expShape), 0,
"BATCHNORM_BP CUDNN op: wrong shape of mean array, expected is %s, but got %s instead !",
ShapeUtils::shapeAsString(expShape).c_str(), ShapeUtils::shapeAsString(mean).c_str());
REQUIRE_TRUE(variance->isSameShape(expShape), 0,
"BATCHNORM_BP CUDNN op: wrong shape of variance array, expected is %s, but got %s instead !",
ShapeUtils::shapeAsString(expShape).c_str(), ShapeUtils::shapeAsString(variance).c_str());
if (gamma)
REQUIRE_TRUE(gamma->isSameShape(expShape), 0,
"BATCHNORM_BP CUDNN op: wrong shape of gamma array, expected is %s, but got %s instead !",
ShapeUtils::shapeAsString(expShape).c_str(), ShapeUtils::shapeAsString(gamma).c_str());
if (beta)
REQUIRE_TRUE(beta->isSameShape(expShape), 0,
"BATCHNORM_BP CUDNN op: wrong shape of beta array, expected is %s, but got %s instead !",
ShapeUtils::shapeAsString(expShape).c_str(), ShapeUtils::shapeAsString(beta).c_str());
REQUIRE_TRUE(input->isSameShape(gradO), 0,
"BATCHNORM_BP CUDNN op: wrong shape of output gradients array, expected is %s, but got %s instead !",
ShapeUtils::shapeAsString(input).c_str(), ShapeUtils::shapeAsString(gradO).c_str());
// types of all input arrays should be the same (except gradO)
for (int i = 1; i < block.width() - 2; ++i)
REQUIRE_TRUE(INPUT_VARIABLE(0)->dataType() == INPUT_VARIABLE(i)->dataType(), 0,
"BATCHNORM_BP CUDNN op: types of arrays (input, mean, variance, gamma, beta) should be the same !");
// cudnn supports NCHW format only
const bool needPermut = axes.size() == 1 && mean->lengthOf() != input->sizeAt(1);
std::unique_ptr<NDArray> tmpGamma = {}, tmpGradG = {}, tmpGradB = {}, tmpInput = {}, tmpGradI = {}, tmpGradO = {};
if (needPermut) { // if NHWC
std::vector<LongType> perm =
inRank == 4 ? std::vector<LongType>({0, 3, 1, 2}) : std::vector<LongType>({0, 4, 1, 2, 3}); // NHWC -> NCHW
tmpInput.reset(new NDArray(input->permute(perm)));
tmpGradO.reset(new NDArray(gradO->permute(perm)));
tmpGradI.reset(new NDArray(gradI->permute(perm)));
input = tmpInput.get();
gradO = tmpGradO.get();
gradI = tmpGradI.get();
}
// cudnn requires gamma, gradG, gradB to be non-nullptr
if (!applyScale) {
tmpGamma.reset(new NDArray(mean));
tmpGradG.reset(new NDArray(mean));
gamma = tmpGamma.get();
gradG = tmpGradG.get();
*gamma = 1;
}
if (!applyOffset) {
tmpGradB.reset(new NDArray(mean));
gradB = tmpGradB.get();
}
// calculations
batchnormBpCUDNN(block.launchContext(), input, mean, variance, gamma, gradO, gradI, gradG, gradB, epsilon,
axes.size() == 1);
*gradM = 0; // put zeros so far
*gradV = 0; // put zeros so far
return Status::OK;
}
PLATFORM_CHECK(batchnorm_bp, ENGINE_CUDA) {
NDArray* input = INPUT_VARIABLE(0);
NDArray* mean = INPUT_VARIABLE(1);
NDArray* variance = INPUT_VARIABLE(2);
NDArray* gamma = nullptr;
NDArray* beta = nullptr;
NDArray* gradO = INPUT_VARIABLE(block.width() - 1); // next epsilon
NDArray* gradI = OUTPUT_VARIABLE(0);
NDArray* gradM = OUTPUT_VARIABLE(1);
NDArray* gradV = OUTPUT_VARIABLE(2);
NDArray* gradG = nullptr;
NDArray* gradB = nullptr;
const int numOfIntArgs = block.getIArguments()->size();
const int xRank = input->rankOf();
// *********************************** //
// get axes args to normalize input array over
std::vector<int> axes;
if (numOfIntArgs > 2)
for (int i = 2; i < numOfIntArgs; ++i) axes.push_back(INT_ARG(i));
else
axes.push_back(xRank - 1); // default dimension to reduce along is last dimension
Requirements req("CUDNN BATCHNORM_BP OP");
req.expectIn(makeInfoVariable(xRank, RANK_MSG_INPUT0), {4, 5}) &&
req.expectIn(makeInfoVariable(input->dataType(), TYPE_MSG_INPUT0),
{HALF, FLOAT32, DOUBLE}) &&
req.expectIn(makeInfoVariable(axes.size(), "axes.size()"), {1, 3, 4}) &&
req.expect(
makeShapeInfoVariable(mean, SHAPE_MSG_INPUT1), makeShapeInfoVariable(variance, SHAPE_MSG_INPUT2),
[](const decltype(mean)& l, const decltype(variance)& r) {
return shape::haveSameShapeAndStrides(l->shapeInfo(), r->shapeInfo());
},
EXPECTED_EQ_MSG);
if (gamma) {
req.expect(
makeShapeInfoVariable(gamma, SHAPE_MSG_INPUT_ "#gamma"), makeShapeInfoVariable(mean, SHAPE_MSG_INPUT1),
[](const decltype(gamma)& l, const decltype(mean)& r) {
return shape::haveSameShapeAndStrides(l->shapeInfo(), r->shapeInfo());
},
EXPECTED_EQ_MSG);
}
if (gradG) {
req.expect(
makeShapeInfoVariable(gradG, SHAPE_MSG_INPUT_ "#gradG"), makeShapeInfoVariable(mean, SHAPE_MSG_INPUT1),
[](const decltype(gradG)& l, const decltype(mean)& r) {
return shape::haveSameShapeAndStrides(l->shapeInfo(), r->shapeInfo());
},
EXPECTED_EQ_MSG);
}
if (gradB) {
req.expect(
makeShapeInfoVariable(gradB, SHAPE_MSG_INPUT_ "#gradB"), makeShapeInfoVariable(mean, SHAPE_MSG_INPUT1),
[](const decltype(gradB)& l, const decltype(mean)& r) {
return shape::haveSameShapeAndStrides(l->shapeInfo(), r->shapeInfo());
},
EXPECTED_EQ_MSG);
}
if (axes.size() == 1) {
// isFormatGood = mean->lengthOf() == input->sizeAt(1) || mean->lengthOf() == input->sizeAt(-1); // mean [C]
req.expectIn(makeInfoVariable(mean->lengthOf(), LENGTH_MSG_INPUT1), {-1, 1});
} else {
auto inputShapeModif = input->getShapeAsVector(); // [dim0,dim1,dim2,dim3] 4D or [dim0,dim1,dim2,dim3,dim4]
inputShapeModif[0] = 1;
// isFormatGood = mean->isSameShape(inputShapeModif); // mean [1,dim1,dim2,dim3] 4D or
// [1,dim1,dim2,dim3,dim4]
req.expect(
makeShapeInfoVariable(mean, SHAPE_MSG_INPUT1),
makeShapeInfoVariable(inputShapeModif, SHAPE_MSG_INPUT_ "#expect"),
[](const decltype(mean)& l, const decltype(inputShapeModif)& r) { return l->isSameShape(r); }, EXPECTED_EQ_MSG);
}
req.logTheSuccess();
return req;
}
} // namespace platforms
} // namespace ops
} // namespace sd
@@ -0,0 +1,529 @@
/* ******************************************************************************
*
*
* 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 <ops/declarable/helpers/convolutions.h>
#include "cudnnUtils.h"
namespace sd {
namespace ops {
namespace platforms {
//////////////////////////////////////////////////////////////////////////
static void conv2dCUDNN(const LaunchContext* context, NDArray* input, NDArray* weights, NDArray* bias,
NDArray* output, const int kH, const LongType kW, const LongType sH, const LongType sW, const LongType pH,
const LongType pW, const LongType dH, const LongType dW, const int paddingMode, const bool isNCHW,
const int wFormat) {
// cudnn support only two formats for weights {oC,iC,kH,kW} and {oC,kH,kW,iC}
LongType bS, iC, iH, iW, oC, oH,
oW; // batch size, input channels, input height/width, output channels, output height/width;
LongType indIOioC, indIiH, indWoC, indWiC, indWkH, indOoH; // corresponding indexes
ConvolutionUtils::getSizesAndIndexesConv2d(isNCHW, wFormat, *input, *output, bS, iC, iH, iW, oC, oH, oW, indIOioC,
indIiH, indWiC, indWoC, indWkH, indOoH);
auto handle = reinterpret_cast<cudnnHandle_t*>(context->getCuDnnHandle());
CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnSetStream), cudnnSetStream(*handle, *context->getCudaStream()));
cudnnTensorFormat_t format = isNCHW ? CUDNN_TENSOR_NCHW : CUDNN_TENSOR_NHWC;
cudnnTensorFormat_t formatW = 0 == wFormat ? format : (1 == wFormat ? CUDNN_TENSOR_NCHW : CUDNN_TENSOR_NHWC);
// input descriptor
CudnnTensor x;
if (input->ordering() == 'c')
x.set4D(format, cudnnDataType(input->dataType()), bS, iC, iH, iW);
else
x.set4DEx(cudnnDataType(input->dataType()), bS, iC, iH, iW, input->strideAt(0), input->strideAt(indIOioC),
input->strideAt(indIiH), input->strideAt(indIiH + 1));
// weights descriptor
FilterDesc w;
w.set4D(cudnnDataType(weights->dataType()), formatW, oC, iC, kH, kW);
// output descriptor
CudnnTensor z;
if (output->ordering() == 'c')
z.set4D(format, cudnnDataType(output->dataType()), bS, oC, oH, oW);
else
z.set4DEx(cudnnDataType(output->dataType()), bS, oC, oH, oW, output->strideAt(0), output->strideAt(indIOioC),
output->strideAt(indOoH), output->strideAt(indOoH + 1));
// description of convolution
ConvolutionDesc conv;
conv.set2D(pH, pW, sH, sW, dH, dW, CUDNN_CROSS_CORRELATION, cudnnDataType(output->dataType()));
// algorithm description
cudnnConvolutionFwdAlgo_t algo;
cudnnConvolutionFwdAlgoPerf_t algoPerf;
int count = 0;
// err = cudnnGetConvolutionForwardAlgorithm(*handle, x, w, conv, z, CUDNN_CONVOLUTION_FWD_PREFER_FASTEST, 0, &algo);
CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnFindConvolutionForwardAlgorithm),
cudnnFindConvolutionForwardAlgorithm(*handle, x, w, conv, z, 1, &count, &algoPerf));
if (count == 0)
throw cuda_exception::build("conv2dCUDNN: cudnnGetConvolutionForwardAlgorithm failed as the count is 0", 0);
algo = algoPerf.algo;
PointersManager manager(context, __func__);
// allocate auxiliary device memory, abbreviation ws means workspace
size_t wsSize;
CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnGetConvolutionForwardWorkspaceSize),
cudnnGetConvolutionForwardWorkspaceSize(*handle, x, w, conv, z, algo, &wsSize));
void* wsData = manager.allocateDevMem(wsSize);
// provide scaling parameters
const float alpha32(1), beta32(0);
const double alpha64(1), beta64(0);
const void* alpha =
output->sizeOfT() <= 4 ? reinterpret_cast<const void*>(&alpha32) : reinterpret_cast<const void*>(&alpha64);
const void* beta =
output->sizeOfT() <= 4 ? reinterpret_cast<const void*>(&beta32) : reinterpret_cast<const void*>(&beta64);
NDArray::prepareSpecialUse({output}, {input, weights, bias});
// run calculation
CHECK_CUDNN_FAILURE_MSG(
STRINGIZE(cudnnConvolutionForward),
cudnnConvolutionForward(*handle, alpha, x, input->specialBuffer(), w, weights->specialBuffer(), conv, algo,
wsData, wsSize, beta, z, output->specialBuffer()));
// add bias if it is present
if (bias != nullptr) {
CudnnTensor b;
b.set4D(CUDNN_TENSOR_NCHW, cudnnDataType(bias->dataType()), 1, oC, 1, 1);
CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnAddTensor), cudnnAddTensor(*handle, alpha, b, bias->specialBuffer(), alpha,
z, output->specialBuffer()));
}
NDArray::registerSpecialUse({output}, {input, weights, bias});
}
//////////////////////////////////////////////////////////////////////////
static void conv2dBpCUDNN(const LaunchContext* context, NDArray* input, NDArray* weights,
NDArray* gradO, NDArray* gradI, NDArray* gradW, NDArray* gradB, const LongType kH,
const LongType kW, const LongType sH, const LongType sW, const LongType pH, const LongType pW, const LongType dH,
const LongType dW, const LongType paddingMode, const bool isNCHW, const int wFormat) {
LongType bS, iC, iH, iW, oC, oH,
oW; // batch size, input channels, input height/width, output channels, output height/width;
LongType indIOioC, indIiH, indWoC, indWiC, indWkH, indOoH; // corresponding indexes
ConvolutionUtils::getSizesAndIndexesConv2d(isNCHW, wFormat, *input, *gradO, bS, iC, iH, iW, oC, oH, oW, indIOioC,
indIiH, indWiC, indWoC, indWkH, indOoH);
auto handle = reinterpret_cast<cudnnHandle_t*>(context->getCuDnnHandle());
CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnSetStream), cudnnSetStream(*handle, *context->getCudaStream()));
cudnnTensorFormat_t format = isNCHW ? CUDNN_TENSOR_NCHW : CUDNN_TENSOR_NHWC;
cudnnTensorFormat_t formatW = 0 == wFormat ? format : (1 == wFormat ? CUDNN_TENSOR_NCHW : CUDNN_TENSOR_NHWC);
PointersManager manager(context, __func__);
// input descriptor, gradO descriptor, gradI descriptor
CudnnTensor x, dz, dx;
if (input->ordering() == 'c')
x.set4D(format, cudnnDataType(input->dataType()), bS, iC, iH, iW);
else
x.set4DEx(cudnnDataType(input->dataType()), bS, iC, iH, iW, input->strideAt(0), input->strideAt(indIOioC),
input->strideAt(indIiH), input->strideAt(indIiH + 1));
if (gradO->ordering() == 'c')
dz.set4D(format, cudnnDataType(gradO->dataType()), bS, oC, oH, oW);
else
dz.set4DEx(cudnnDataType(gradO->dataType()), bS, oC, oH, oW, gradO->strideAt(0), gradO->strideAt(indIOioC),
gradO->strideAt(indOoH), gradO->strideAt(indOoH + 1));
if (gradI->ordering() == 'c')
dx.set4D(format, cudnnDataType(gradI->dataType()), bS, iC, iH, iW);
else
dx.set4DEx(cudnnDataType(gradI->dataType()), bS, iC, iH, iW, gradI->strideAt(0), gradI->strideAt(indIOioC),
gradI->strideAt(indIiH), gradI->strideAt(indIiH + 1));
// gradW descriptor
FilterDesc dw;
dw.set4D(cudnnDataType(gradW->dataType()), formatW, oC, iC, kH, kW);
// description of convolution
ConvolutionDesc conv;
conv.set2D(pH, pW, sH, sW, dH, dW, CUDNN_CROSS_CORRELATION, cudnnDataType(gradO->dataType()));
// gradW algorithm description
cudnnConvolutionBwdFilterAlgo_t algoGradW;
cudnnConvolutionBwdFilterAlgoPerf_t algoGradWPerf;
int count = 0;
CHECK_CUDNN_FAILURE_MSG(
STRINGIZE(cudnnFindConvolutionBackwardFilterAlgorithm),
cudnnFindConvolutionBackwardFilterAlgorithm(*handle, x, dz, conv, dw, 1, &count, &algoGradWPerf));
if (count == 0)
throw cuda_exception::build(
"conv2dBpCUDNN: cudnnGetConvolutionBackwardFilterAlgorithm failed as the count is 0", 0);
algoGradW = algoGradWPerf.algo;
// gradI algorithm description
cudnnConvolutionBwdDataAlgo_t algoGradI;
cudnnConvolutionBwdDataAlgoPerf_t algoGradIPerf;
CHECK_CUDNN_FAILURE_MSG(
STRINGIZE(cudnnFindConvolutionBackwardDataAlgorithm),
cudnnFindConvolutionBackwardDataAlgorithm(*handle, dw, dz, conv, x, 1, &count, &algoGradIPerf));
if (count == 0)
throw cuda_exception::build("conv2dBpCUDNN: cudnnGetConvolutionBackwardDataAlgorithm failed as the count is 0",
0);
algoGradI = algoGradIPerf.algo;
// allocate auxiliary device memory for gradW calculation, abbreviation ws means workspace
size_t wsGradWSize;
CHECK_CUDNN_FAILURE_MSG(
STRINGIZE(cudnnGetConvolutionBackwardFilterWorkspaceSize),
cudnnGetConvolutionBackwardFilterWorkspaceSize(*handle, x, dz, conv, dw, algoGradW, &wsGradWSize));
void* wsGradWData = manager.allocateDevMem(wsGradWSize);
// allocate auxiliary device memory for gradI calculation, abbreviation ws means workspace
size_t wsGradISize;
CHECK_CUDNN_FAILURE_MSG(
STRINGIZE(cudnnGetConvolutionBackwardDataWorkspaceSize),
cudnnGetConvolutionBackwardDataWorkspaceSize(*handle, dw, dz, conv, dx, algoGradI, &wsGradISize));
void* wsGradIData = manager.allocateDevMem(wsGradISize);
// provide scaling parameters
const float alpha32(1), beta32(0);
const double alpha64(1), beta64(0);
const void* alpha =
gradO->sizeOfT() <= 4 ? reinterpret_cast<const void*>(&alpha32) : reinterpret_cast<const void*>(&alpha64);
const void* beta =
gradO->sizeOfT() <= 4 ? reinterpret_cast<const void*>(&beta32) : reinterpret_cast<const void*>(&beta64);
NDArray::prepareSpecialUse({gradI, gradW, gradB}, {input, weights, gradO});
// run calculation for gradB (if not nullptr)
if (gradB != nullptr) {
CudnnTensor db;
db.set4D(CUDNN_TENSOR_NCHW, cudnnDataType(gradB->dataType()), 1, oC, 1, 1);
CHECK_CUDNN_FAILURE_MSG(
STRINGIZE(cudnnConvolutionBackwardBias),
cudnnConvolutionBackwardBias(*handle, alpha, dz, gradO->specialBuffer(), beta, db, gradB->specialBuffer()));
}
// run calculation for gradW
CHECK_CUDNN_FAILURE_MSG(
STRINGIZE(cudnnConvolutionBackwardFilter),
cudnnConvolutionBackwardFilter(*handle, alpha, x, input->specialBuffer(), dz, gradO->specialBuffer(), conv,
algoGradW, wsGradWData, wsGradWSize, beta, dw, gradW->specialBuffer()));
// run calculation for gradI
CHECK_CUDNN_FAILURE_MSG(
STRINGIZE(cudnnConvolutionBackwardData),
cudnnConvolutionBackwardData(*handle, alpha, dw, weights->specialBuffer(), dz, gradO->specialBuffer(), conv,
algoGradI, wsGradIData, wsGradISize, beta, dx, gradI->specialBuffer()));
NDArray::registerSpecialUse({gradI, gradW, gradB}, {input, weights, gradO});
}
//////////////////////////////////////////////////////////////////////////
PLATFORM_IMPL(conv2d, ENGINE_CUDA) {
auto input = INPUT_VARIABLE(0); // [bS, iH, iW, iC] (NHWC) or [bS, iC, iH, iW] (NCHW)
auto weights = INPUT_VARIABLE(1); // [kH, kW, iC, oC], [oC, iC, kH, kW], [oC, kH, kW, iC]
auto bias = block.width() > 2 ? INPUT_VARIABLE(2) : nullptr; // [oC]
auto output = OUTPUT_VARIABLE(0); // [bS, oH, oW, oC] (NHWC) or [bS, oC, oH, oW] (NCHW)
LongType sH = INT_ARG(2); // strides height
LongType sW = INT_ARG(3); // strides width
LongType pH = INT_ARG(4); // paddings height
LongType pW = INT_ARG(5); // paddings width
LongType dH = INT_ARG(6); // dilations height
LongType dW = INT_ARG(7); // dilations width
int paddingMode = INT_ARG(8); // 0-VALID, 1-SAME
bool isNCHW = block.getIArguments()->size() > 9 ? !INT_ARG(9) : 1; // INT_ARG(9): 0-NCHW, 1-NHWC
int wFormat = block.getIArguments()->size() > 10
? INT_ARG(10)
: 0; // 0 - [kH, kW, iC, oC], 1 - [oC, iC, kH, kW], 2 - [oC, kH, kW, iC]
LongType kH = INT_ARG(0) > 0 ? INT_ARG(0) : static_cast<LongType>(weights->sizeAt(0)); // filter(kernel) height
LongType kW = INT_ARG(1) > 0 ? INT_ARG(1) : static_cast<LongType>(weights->sizeAt(1)); // filter(kernel) width
REQUIRE_TRUE(input->rankOf() == 4, 0,
"CUSTOM CONV2D CUDNN OP: rank of input array must be equal to 4, but got %i instead !", input->rankOf());
REQUIRE_TRUE(weights->rankOf() == 4, 0,
"CUSTOM CONV2D CUDNN OP: rank of weights array must be equal to 4, but got %i instead !",
weights->rankOf());
LongType bS, iC, iH, iW, oC, oH,
oW; // batch size, input channels, input height/width, output channels, output height/width;
LongType indIOioC, indIiH, indWoC, indWiC, indWkH, indOoH; // corresponding indexes
ConvolutionUtils::getSizesAndIndexesConv2d(isNCHW, wFormat, *input, *output, bS, iC, iH, iW, oC, oH, oW, indIOioC,
indIiH, indWiC, indWoC, indWkH, indOoH);
ConvolutionUtils::calcPadding2D(pH, pW, oH, oW, iH, iW, kH, kW, sH, sW, dH, dW, paddingMode);
std::vector<LongType> expectedWeightsShape = ConvolutionUtils::expectWeightsShape(wFormat, kH, kW, iC, oC);
REQUIRE_TRUE(weights->isSameShape(expectedWeightsShape), 0,
"CUSTOM CONV2D CUDNN OP: wrong shape of weights array, expected is %s, but got %s instead !",
ShapeUtils::shapeAsString(expectedWeightsShape).c_str(), ShapeUtils::shapeAsString(weights).c_str());
if (bias) {
REQUIRE_TRUE(bias->rankOf() <= 2 && oC == bias->lengthOf(), 0,
"CUSTOM CONV2D CUDNN OP: wrong shape of array with biases, expected rank, length: <=2, %i, but got "
"%i, %i instead !",
oC, bias->rankOf(), bias->lengthOf());
REQUIRE_TRUE((bias->rankOf() == 1 && bias->strideAt(0) == 1) ||
(bias->rankOf() == 2 && bias->sizeAt(0) == 1 && bias->strideAt(1) == 1) ||
(bias->rankOf() == 2 && bias->sizeAt(1) == 1 && bias->strideAt(0) == 1),
0, "CUSTOM CONV2D CUDNN OP: bias array should be contiguous in memory !");
}
std::unique_ptr<NDArray> tmpWeight = {}, tmpInput = {};
NDArray* newWeights = weights; // cudnn support only two formats {oC,iC,kH,kW} and {oC,kH,kW,iC}
if (0 == wFormat) {
// Create named vectors as lvalues
std::vector<LongType> nchwShape = {oC, iC, kH, kW};
std::vector<LongType> nhwcShape = {oC, kH, kW, iC};
// Use the appropriate one for the weight reset
tmpWeight.reset(
new NDArray(weights->ordering(),
isNCHW ? nchwShape : nhwcShape,
weights->dataType(), weights->getContext()));
newWeights = tmpWeight.get();
// Create named vectors as lvalues
std::vector<LongType> nchwDims = {3, 2, 0, 1};
std::vector<LongType> nhwcDims = {3, 0, 1, 2};
// Use the appropriate one in the call
NDArray assign = weights->permute(
isNCHW ? nchwDims : nhwcDims,
true, // copyToNewBuff
true); // resetStrides
newWeights->assign(&assign); // (kH, kW, iC, oC --> oC, iC, kH, kW) or (kH, kW, iC, oC --> oC, kH, kW, iC)
}
if (paddingMode == 1) { // in same paddingMode cudnn doesn't support asymmetric left/right top/bottopm paddings
auto ret = checkConv2dCUDNNPadAsymmetric(input, nullptr, iH, iW, oH, oW, kH, kW, sH, sW, pH, pW, dH, dW, isNCHW);
tmpInput = std::move(std::get<0>(ret)); // prolong life
if (tmpInput) input = tmpInput.get();
}
conv2dCUDNN(block.launchContext(), input, newWeights, bias, output, kH, kW, sH, sW, pH, pW, dH, dW, paddingMode,
isNCHW, wFormat);
return Status::OK;
}
//////////////////////////////////////////////////////////////////////////
PLATFORM_CHECK(conv2d, ENGINE_CUDA) {
auto input = INPUT_VARIABLE(0); // [bS, iH, iW, iC] (NHWC) or [bS, iC, iH, iW] (NCHW)
auto weights = INPUT_VARIABLE(1); // [kH, kW, iC, oC] always
auto bias = block.width() > 2 ? INPUT_VARIABLE(2) : nullptr; // [oC]
const int paddingMode = INT_ARG(8); // 0-VALID, 1-SAME, 2-CAUSAL
const bool badInputType = input->dataType() != DOUBLE && input->dataType() != FLOAT32 &&
input->dataType() != HALF;
const bool badWeightsType = weights->dataType() != DOUBLE && weights->dataType() != FLOAT32 &&
weights->dataType() != HALF;
const bool badBiasType = bias == nullptr
? false
: (bias->dataType() != DOUBLE && bias->dataType() != FLOAT32 &&
bias->dataType() != HALF);
return paddingMode != 2 && !badInputType && !badWeightsType && !badBiasType;
Requirements req("CUDNN CONV2d OP");
req.expectNotEq(makeInfoVariable(paddingMode, "paddingMode"), 2) &&
req.expectIn(makeInfoVariable(input->dataType(), TYPE_MSG_INPUT0),
{HALF, FLOAT32, DOUBLE}) &&
req.expectIn(makeInfoVariable(weights->dataType(), TYPE_MSG_INPUT1),
{HALF, FLOAT32, DOUBLE});
if (bias) {
req.expectIn(makeInfoVariable(bias->dataType(), TYPE_MSG_INPUT_ "#bias"),
{HALF, FLOAT32, DOUBLE});
}
req.logTheSuccess();
return req;
}
//////////////////////////////////////////////////////////////////////////
PLATFORM_IMPL(conv2d_bp, ENGINE_CUDA) {
auto input = INPUT_VARIABLE(0); // [bS, iH, iW, iC] (NHWC) or [bS, iC, iH, iW] (NCHW)
auto weights = INPUT_VARIABLE(1); // [kH, kW, iC, oC], [oC, iC, kH, kW], [oC, kH, kW, iC]
auto bias = block.width() > 3 ? INPUT_VARIABLE(2) : nullptr; // [oC]
auto gradO = block.width() > 3
? INPUT_VARIABLE(3)
: INPUT_VARIABLE(2); // [bS, oH, oW, oC] (NHWC) or [bS, oC, oH, oW] (NCHW), epsilon_next
auto gradI = OUTPUT_VARIABLE(0); // [bS, iH, iW, iC] (NHWC) or [bS, iC, iH, iW] (NCHW), epsilon
auto gradW = OUTPUT_VARIABLE(1); // [kH, kW, iC, oC], [oC, iC, kH, kW], [oC, kH, kW, iC]
auto gradB = block.width() > 3 ? OUTPUT_VARIABLE(2) : nullptr; // [oC]
LongType kH = INT_ARG(0); // filter(kernel) height
LongType kW = INT_ARG(1); // filter(kernel) width
LongType sH = INT_ARG(2); // strides height
LongType sW = INT_ARG(3); // strides width
LongType pH = INT_ARG(4); // paddings height
LongType pW = INT_ARG(5); // paddings width
LongType dH = INT_ARG(6); // dilations height
LongType dW = INT_ARG(7); // dilations width
int paddingMode = INT_ARG(8); // 0-VALID, 1-SAME
int isNCHW = block.getIArguments()->size() > 9 ? !INT_ARG(9) : 1; // INT_ARG(9): 0-NCHW, 1-NHWC
int wFormat = block.getIArguments()->size() > 10
? INT_ARG(10)
: 0; // 0 - [kH, kW, iC, oC], 1 - [oC, iC, kH, kW], 2 - [oC, kH, kW, iC]
REQUIRE_TRUE(input->rankOf() == 4, 0,
"CUSTOM CONV2D_BP CUDNN OP: rank of input array must be equal to 4, but got %i instead !",
input->rankOf());
REQUIRE_TRUE(weights->rankOf() == 4, 0,
"CUSTOM CONV2D_BP CUDNN OP: rank of weights array must be equal to 4, but got %i instead !",
weights->rankOf());
REQUIRE_TRUE(gradO->rankOf() == 4, 0,
"CUSTOM CONV2D_BP CUDNN OP: rank of output's gradients (next epsilon) array must be equal to 4, but got "
"%i instead !",
gradO->rankOf());
LongType bS, iC, iH, iW, oC, oH,
oW; // batch size, input channels, input height/width, output channels, output height/width;
LongType indIOioC, indIiH, indWoC, indWiC, indWkH, indOoH; // corresponding indexes
ConvolutionUtils::getSizesAndIndexesConv2d(isNCHW, wFormat, *input, *gradO, bS, iC, iH, iW, oC, oH, oW, indIOioC,
indIiH, indWiC, indWoC, indWkH, indOoH);
LongType trueoH, trueoW; // true output height, width
ConvolutionUtils::calcOutSizePool2D(trueoH, trueoW, kH, kW, sH, sW, pH, pW, dH, dW, iH, iW, paddingMode);
ConvolutionUtils::calcPadding2D(pH, pW, oH, oW, iH, iW, kH, kW, sH, sW, dH, dW, paddingMode);
std::vector<LongType> expectedGradOShape =
ShapeUtils::composeShapeUsingDimsAndIdx({bS, oC, trueoH, trueoW, 0, indIOioC, indOoH, indOoH + 1});
std::vector<LongType> expectedWeightsShape = ConvolutionUtils::expectWeightsShape(wFormat, kH, kW, iC, oC);
REQUIRE_TRUE(gradO->isSameShape(expectedGradOShape), 0,
"CUSTOM CONV2D_BP CUDNN OP: wrong shape of output gradients (next epsilon) array, expected is %s, but "
"got %s instead !",
ShapeUtils::shapeAsString(expectedGradOShape).c_str(), ShapeUtils::shapeAsString(gradO).c_str());
REQUIRE_TRUE(weights->isSameShape(expectedWeightsShape), 0,
"CUSTOM CONV2D_BP CUDNN OP: wrong shape of weights array, expected is %s, but got %s instead !",
ShapeUtils::shapeAsString(expectedWeightsShape).c_str(), ShapeUtils::shapeAsString(weights).c_str());
if (bias)
REQUIRE_TRUE(bias->rankOf() <= 2 && oC == bias->lengthOf(), 0,
"CUSTOM CONV2D_BP CUDNN OP: wrong shape of array with biases, expected rank, length: <=2, %i, but got "
"%i, %i instead !",
oC, bias->rankOf(), bias->lengthOf());
std::unique_ptr<NDArray> tmpGradI = {}, tmpInput = {}, tmpWeights = {}, tmpGradW = {};
NDArray *newWeights = weights, *newGradW = gradW; // cudnn support only two formats {oC,iC,kH,kW} and {oC,kH,kW,iC}
if (0 == wFormat) {
// Create named vectors as lvalues
std::vector<LongType> nchwGradShape = {oC, iC, kH, kW};
std::vector<LongType> nhwcGradShape = {oC, kH, kW, iC};
// Use the appropriate one for the gradW reset
tmpGradW.reset(
new NDArray(gradW->ordering(),
isNCHW ? nchwGradShape : nhwcGradShape,
gradW->dataType(), gradW->getContext()));
// Use the same vectors for the weights reset
tmpWeights.reset(
new NDArray(weights->ordering(),
isNCHW ? nchwGradShape : nhwcGradShape,
weights->dataType(), weights->getContext()));
newGradW = tmpGradW.get();
newWeights = tmpWeights.get();
// Create named vectors as lvalues
std::vector<LongType> nchwDims = {3, 2, 0, 1};
std::vector<LongType> nhwcDims = {3, 0, 1, 2};
NDArray assign = weights->permute(
isNCHW ? nchwDims : nhwcDims,
true, // copyToNewBuff
true);
// Use the appropriate one in the call
newWeights->assign(&assign);
}
NDArray* newInput = input;
NDArray* newGradI = gradI;
if (paddingMode == 1) { // in same paddingMode cudnn doesn't support asymmetric left/right top/bottopm paddings
auto ret = checkConv2dCUDNNPadAsymmetric(input, gradI, iH, iW, oH, oW, kH, kW, sH, sW, pH, pW, dH, dW, isNCHW);
tmpInput = std::move(std::get<0>(ret));
tmpGradI = std::move(std::get<1>(ret));
if (tmpInput) newInput = tmpInput.get();
if (tmpGradI) newGradI = tmpGradI.get();
}
conv2dBpCUDNN(block.launchContext(), newInput, newWeights, gradO, newGradI, newGradW, gradB, kH, kW, sH, sW, pH, pW,
dH, dW, paddingMode, isNCHW, wFormat);
if (0 == wFormat) {
// Create named vectors as lvalues
std::vector<LongType> nchwPermute = {2, 3, 1, 0};
std::vector<LongType> nhwcPermute = {1, 2, 3, 0};
// Use the appropriate one in the permutei call
newGradW->permutei(
isNCHW ? nchwPermute : nhwcPermute,false,false); // (oC, iC, kH, kW --> kH, kW, iC, oC) or (oC, kH, kW, iC --> kH, kW, iC, oC) iC, oC)
gradW->assign(newGradW);
}
if (newInput != input) {
if (isNCHW) {
NDArray assign = (*newGradI)({0, 0, 0, 0, 0, gradI->sizeAt(2), 0, gradI->sizeAt(3)});
gradI->assign(&assign);
} else {
NDArray assign = (*newGradI)({0, 0, 0, gradI->sizeAt(1), 0, gradI->sizeAt(2), 0, 0});
gradI->assign(&assign);
}
}
return Status::OK;
}
PLATFORM_CHECK(conv2d_bp, ENGINE_CUDA) {
auto input = INPUT_VARIABLE(0); // [bS, iH, iW, iC] (NHWC) or [bS, iC, iH, iW] (NCHW)
auto weights = INPUT_VARIABLE(1); // [kH, kW, iC, oC] always
auto bias = block.width() > 3 ? INPUT_VARIABLE(2) : nullptr; // [oC]
auto gradO = block.width() > 3
? INPUT_VARIABLE(3)
: INPUT_VARIABLE(2); // [bS, oH, oW, oC] (NHWC) or [bS, oC, oH, oW] (NCHW), epsilon_next
const int paddingMode = INT_ARG(8); // 0-VALID, 1-SAME, 2-CAUSAL
const int isNCHW = block.getIArguments()->size() > 9 ? !INT_ARG(9) : 1; // INT_ARG(9): 0-NCHW, 1-NHWC
Requirements req("CUDNN CONV2d_BP OP");
req.expectNotEq(makeInfoVariable(paddingMode, "paddingMode"), 2) &&
req.expectTrue(makeInfoVariable(isNCHW, "isNCHW")) &&
req.expectIn(makeInfoVariable(input->dataType(), TYPE_MSG_INPUT0),
{HALF, FLOAT32, DOUBLE}) &&
req.expectIn(makeInfoVariable(weights->dataType(), TYPE_MSG_INPUT1),
{HALF, FLOAT32, DOUBLE});
if (bias) {
req.expectIn(makeInfoVariable(bias->dataType(), TYPE_MSG_INPUT_ "#bias"),
{HALF, FLOAT32, DOUBLE}) &&
req.expectIn(makeInfoVariable(gradO->dataType(), TYPE_MSG_INPUT3),
{HALF, FLOAT32, DOUBLE});
} else {
req.expectIn(makeInfoVariable(gradO->dataType(), TYPE_MSG_INPUT2),
{HALF, FLOAT32, DOUBLE});
}
req.logTheSuccess();
return req;
}
} // namespace platforms
} // namespace ops
} // namespace sd
@@ -0,0 +1,543 @@
/* ******************************************************************************
*
*
* 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 <ops/declarable/helpers/convolutions.h>
#include "cudnnUtils.h"
namespace sd {
namespace ops {
namespace platforms {
//////////////////////////////////////////////////////////////////////////
static void conv3dCUDNN(const LaunchContext* context, NDArray* input, NDArray* weights, NDArray* bias,
NDArray* output, const LongType kD, const LongType kH, const LongType kW, const LongType sD, const LongType sH,
const LongType sW, const LongType pD, const LongType pH, const LongType pW, const LongType dD, const LongType dH,
const LongType dW, const int paddingMode, const bool isNCDHW, const int wFormat) {
// cudnn support only one format for weights {oC,iC,kD,kH,kW}
const int numDims = 5;
LongType bS, iC, iD, iH, iW, oC, oD, oH,
oW; // batch size, input channels, input depth/height/width, output channels, output depth/height/width;
LongType indIOioC, indIOioD, indWoC, indWiC, indWkD; // corresponding indexes
ConvolutionUtils::getSizesAndIndexesConv3d(isNCDHW, wFormat, *input, *output, bS, iC, iD, iH, iW, oC, oD, oH, oW,
indIOioC, indIOioD, indWiC, indWoC, indWkD);
auto handle = reinterpret_cast<cudnnHandle_t*>(context->getCuDnnHandle());
CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnSetStream), cudnnSetStream(*handle, *context->getCudaStream()));
const std::vector<int> pads = {static_cast<int>(pD), static_cast<int>(pH), static_cast<int>(pW)};
const std::vector<int> filtStrides = {static_cast<int>(sD), static_cast<int>(sH), static_cast<int>(sW)};
const std::vector<int> dilations = {static_cast<int>(dD), static_cast<int>(dH), static_cast<int>(dW)};
const std::vector<int> xShape = {static_cast<int>(bS), static_cast<int>(iC), static_cast<int>(iD), static_cast<int>(iH), static_cast<int>(iW)};
const std::vector<int> zShape = {static_cast<int>(bS), static_cast<int>(oC), static_cast<int>(oD), static_cast<int>(oH), static_cast<int>(oW)};
const std::vector<int> wShape = {static_cast<int>(oC), static_cast<int>(iC), static_cast<int>(kD), static_cast<int>(kH), static_cast<int>(kW)};
const std::vector<int> bShape = {1, static_cast<int>(oC), 1, 1, 1};
const std::vector<int> xStrides = {static_cast<int>(input->strideAt(0)), static_cast<int>(input->strideAt(1)), static_cast<int>(input->strideAt(2)),
static_cast<int>(input->strideAt(3)), static_cast<int>(input->strideAt(4))};
const std::vector<int> zStrides = {static_cast<int>(output->strideAt(0)), static_cast<int>(output->strideAt(1)), static_cast<int>(output->strideAt(2)),
static_cast<int>(output->strideAt(3)), static_cast<int>(output->strideAt(4))};
cudnnTensorFormat_t format = isNCDHW ? CUDNN_TENSOR_NCHW : CUDNN_TENSOR_NHWC;
PointersManager manager(context, __func__);
// input descriptor
CudnnTensor x;
x.set(cudnnDataType(input->dataType()), numDims, xShape.data(), xStrides.data());
// weights descriptor
FilterDesc w;
w.set(cudnnDataType(weights->dataType()), CUDNN_TENSOR_NCHW, numDims, wShape.data());
// output descriptor
CudnnTensor z;
z.set(cudnnDataType(output->dataType()), numDims, zShape.data(), zStrides.data());
// description of convolution
ConvolutionDesc conv;
conv.set(numDims - 2, pads.data(), filtStrides.data(), dilations.data(), CUDNN_CROSS_CORRELATION,
cudnnDataType(output->dataType()));
// algorithm description
cudnnConvolutionFwdAlgo_t algo;
cudnnConvolutionFwdAlgoPerf_t algoPerf;
int count = 0;
CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnFindConvolutionForwardAlgorithm),
cudnnFindConvolutionForwardAlgorithm(*handle, x, w, conv, z, 1, &count, &algoPerf));
if (count == 0)
throw cuda_exception::build("conv3dCUDNN: cudnnGetConvolutionForwardAlgorithm failed as the count is 0", 0);
algo = algoPerf.algo;
// allocate auxiliary device memory, abbreviation ws means workspace
size_t wsSize;
CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnGetConvolutionForwardWorkspaceSize),
cudnnGetConvolutionForwardWorkspaceSize(*handle, x, w, conv, z, algo, &wsSize));
void* wsData = manager.allocateDevMem(wsSize);
// provide scaling parameters
const float alpha32(1), beta32(0);
const double alpha64(1), beta64(0);
const void* alpha =
output->sizeOfT() <= 4 ? reinterpret_cast<const void*>(&alpha32) : reinterpret_cast<const void*>(&alpha64);
const void* beta =
output->sizeOfT() <= 4 ? reinterpret_cast<const void*>(&beta32) : reinterpret_cast<const void*>(&beta64);
NDArray::prepareSpecialUse({output}, {input, weights, bias});
// run calculation
CHECK_CUDNN_FAILURE_MSG(
STRINGIZE(cudnnConvolutionForward),
cudnnConvolutionForward(*handle, alpha, x, input->specialBuffer(), w, weights->specialBuffer(), conv, algo,
wsData, wsSize, beta, z, output->specialBuffer()));
// add bias if it is present
if (bias != nullptr) {
CudnnTensor b;
b.setEx(/*format*/ CUDNN_TENSOR_NCHW, cudnnDataType(bias->dataType()), numDims, bShape.data());
CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnAddTensor), cudnnAddTensor(*handle, alpha, b, bias->specialBuffer(), alpha,
z, output->specialBuffer()));
}
NDArray::registerSpecialUse({output}, {input, weights, bias});
}
//////////////////////////////////////////////////////////////////////////
static void conv3dBpCUDNN(const LaunchContext* context, NDArray* input, NDArray* weights,
NDArray* gradO, NDArray* gradI, NDArray* gradW, NDArray* gradB, const int kD,
const LongType kH, const LongType kW, const LongType sD, const LongType sH, const LongType sW, const LongType pD,
const LongType pH, const LongType pW, const LongType dD, const LongType dH, const LongType dW, const int paddingMode,
const bool isNCDHW, const int wFormat) {
// cudnn supports only two formats {oC,iC,kD,kH,kW} and {oC,kD,kH,kW,iC} for weights/gradW
const int numDims = 5;
LongType bS, iC, iD, iH, iW, oC, oD, oH,
oW; // batch size, input channels, input depth/height/width, output channels, output depth/height/width;
LongType indIOioC, indIOioD, indWoC, indWiC, indWkD; // corresponding indexes
ConvolutionUtils::getSizesAndIndexesConv3d(isNCDHW, wFormat, *input, *gradO, bS, iC, iD, iH, iW, oC, oD, oH, oW,
indIOioC, indIOioD, indWiC, indWoC, indWkD);
auto handle = reinterpret_cast<cudnnHandle_t*>(context->getCuDnnHandle());
CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnSetStream), cudnnSetStream(*handle, *context->getCudaStream()));
const std::vector<int> pads = {static_cast<int>(pD), static_cast<int>(pH), static_cast<int>(pW)};
const std::vector<int> filtStrides = {static_cast<int>(sD), static_cast<int>(sH), static_cast<int>(sW)};
const std::vector<int> dilations = {static_cast<int>(dD), static_cast<int>(dH), static_cast<int>(dW)};
const std::vector<int> xShape = {static_cast<int>(bS), static_cast<int>(iC), static_cast<int>(iD), static_cast<int>(iH), static_cast<int>(iW)};
const std::vector<int> dzShape = {static_cast<int>(bS), static_cast<int>(oC), static_cast<int>(oD), static_cast<int>(oH), static_cast<int>(oW)};
const std::vector<int> wShape = {static_cast<int>(oC), static_cast<int>(iC), static_cast<int>(kD), static_cast<int>(kH), static_cast<int>(kW)};
const std::vector<int> dbShape = {1, static_cast<int>(isNCDHW ? oC : 1), 1, 1, (int)(isNCDHW ? 1 : oC)};
const std::vector<int> xStrides = {static_cast<int>(input->strideAt(0)), static_cast<int>(input->strideAt(1)), static_cast<int>(input->strideAt(2)),
static_cast<int>(input->strideAt(3)), static_cast<int>(input->strideAt(4))};
const std::vector<int> dxStrides = {static_cast<int>(gradI->strideAt(0)), static_cast<int>(gradI->strideAt(1)), static_cast<int>(gradI->strideAt(2)),
static_cast<int>(gradI->strideAt(3)), static_cast<int>(gradI->strideAt(4))};
const std::vector<int> dzStrides = {static_cast<int>(gradO->strideAt(0)), static_cast<int>(gradO->strideAt(1)), static_cast<int>(gradO->strideAt(2)),
static_cast<int>(gradO->strideAt(3)), static_cast<int>(gradO->strideAt(4))};
cudnnTensorFormat_t format = isNCDHW ? CUDNN_TENSOR_NCHW : CUDNN_TENSOR_NHWC;
cudnnTensorFormat_t formatW = 0 == wFormat ? format : (1 == wFormat ? CUDNN_TENSOR_NCHW : CUDNN_TENSOR_NHWC);
PointersManager manager(context, __func__);
// input descriptor, gradO descriptor, gradI descriptor
CudnnTensor x, dz, dx;
x.set(cudnnDataType(input->dataType()), numDims, xShape.data(), xStrides.data());
dz.set(cudnnDataType(gradO->dataType()), numDims, dzShape.data(), dzStrides.data());
dx.set(cudnnDataType(gradI->dataType()), numDims, xShape.data(), dxStrides.data());
// gradW descriptor
FilterDesc dw;
dw.set(cudnnDataType(gradW->dataType()), formatW, numDims, wShape.data());
// description of convolution
ConvolutionDesc conv;
conv.set(numDims - 2, pads.data(), filtStrides.data(), dilations.data(), CUDNN_CROSS_CORRELATION,
cudnnDataType(gradO->dataType()));
// gradW algorithm description
cudnnConvolutionBwdFilterAlgo_t algoGradW;
cudnnConvolutionBwdFilterAlgoPerf_t algoGradWPerf;
int count = 0;
CHECK_CUDNN_FAILURE_MSG(
STRINGIZE(cudnnFindConvolutionBackwardFilterAlgorithm),
cudnnFindConvolutionBackwardFilterAlgorithm(*handle, x, dz, conv, dw, 1, &count, &algoGradWPerf));
if (count == 0)
throw cuda_exception::build(
"conv3dBpCUDNN: cudnnGetConvolutionBackwardFilterAlgorithm failed as the count is 0", 0);
algoGradW = algoGradWPerf.algo;
// gradI algorithm description
cudnnConvolutionBwdDataAlgo_t algoGradI;
cudnnConvolutionBwdDataAlgoPerf_t algoGradIPerf;
// CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnGetConvolutionBackwardDataAlgorithm),
// cudnnGetConvolutionBackwardDataAlgorithm( *handle, dw, dz, conv, x, CUDNN_CONVOLUTION_BWD_DATA_PREFER_FASTEST, 0,
// &algoGradI));
CHECK_CUDNN_FAILURE_MSG(
STRINGIZE(cudnnFindConvolutionBackwardDataAlgorithm),
cudnnFindConvolutionBackwardDataAlgorithm(*handle, dw, dz, conv, x, 1, &count, &algoGradIPerf));
if (count == 0)
throw cuda_exception::build("conv3dBpCUDNN: cudnnGetConvolutionBackwardDataAlgorithm failed as the count is 0",
0);
algoGradI = algoGradIPerf.algo;
// allocate auxiliary device memory for gradW calculation, abbreviation ws means workspace
size_t wsGradWSize;
CHECK_CUDNN_FAILURE_MSG(
STRINGIZE(cudnnGetConvolutionBackwardFilterWorkspaceSize),
cudnnGetConvolutionBackwardFilterWorkspaceSize(*handle, x, dz, conv, dw, algoGradW, &wsGradWSize));
void* wsGradWData = manager.allocateDevMem(wsGradWSize);
// allocate auxiliary device memory for gradI calculation, abbreviation ws means workspace
size_t wsGradISize;
CHECK_CUDNN_FAILURE_MSG(
STRINGIZE(cudnnGetConvolutionBackwardDataWorkspaceSize),
cudnnGetConvolutionBackwardDataWorkspaceSize(*handle, dw, dz, conv, dx, algoGradI, &wsGradISize));
void* wsGradIData = manager.allocateDevMem(wsGradISize);
// provide scaling parameters
const float alpha32(1), beta32(0);
const double alpha64(1), beta64(0);
const void* alpha =
gradO->sizeOfT() <= 4 ? reinterpret_cast<const void*>(&alpha32) : reinterpret_cast<const void*>(&alpha64);
const void* beta =
gradO->sizeOfT() <= 4 ? reinterpret_cast<const void*>(&beta32) : reinterpret_cast<const void*>(&beta64);
NDArray::prepareSpecialUse({gradI, gradW, gradB}, {input, weights, gradO});
// run calculation for gradB (if not nullptr)
if (gradB != nullptr) {
CudnnTensor db;
db.setEx(format, cudnnDataType(gradB->dataType()), numDims, dbShape.data());
CHECK_CUDNN_FAILURE_MSG(
STRINGIZE(cudnnConvolutionBackwardBias),
cudnnConvolutionBackwardBias(*handle, alpha, dz, gradO->specialBuffer(), beta, db, gradB->specialBuffer()));
}
// run calculation for gradW
CHECK_CUDNN_FAILURE_MSG(
STRINGIZE(cudnnConvolutionBackwardFilter),
cudnnConvolutionBackwardFilter(*handle, alpha, x, input->specialBuffer(), dz, gradO->specialBuffer(), conv,
algoGradW, wsGradWData, wsGradWSize, beta, dw, gradW->specialBuffer()));
// run calculation for gradI
CHECK_CUDNN_FAILURE_MSG(
STRINGIZE(cudnnConvolutionBackwardData),
cudnnConvolutionBackwardData(*handle, alpha, dw, weights->specialBuffer(), dz, gradO->specialBuffer(), conv,
algoGradI, wsGradIData, wsGradISize, beta, dx, gradI->specialBuffer()));
NDArray::registerSpecialUse({gradI, gradW, gradB}, {input, weights, gradO});
}
//////////////////////////////////////////////////////////////////////////
PLATFORM_IMPL(conv3dnew, ENGINE_CUDA) {
auto input = INPUT_VARIABLE(0); // [bS, iD, iH, iW, iC] (NDHWC) or [bS, iC, iD, iH, iW] (NCDHW)
auto weights = INPUT_VARIABLE(1); // [kD, kH, kW, iC, oC], [oC, iC, kD, kH, kW], [oC, kD, kH, kW, iC]
auto bias = block.width() > 2 ? INPUT_VARIABLE(2) : nullptr; // [oC]
auto output = OUTPUT_VARIABLE(0); // [bS, oD, oH, oW, oC] (NDHWC) or [bS, oC, oD, oH, oW] (NCDHW)
REQUIRE_TRUE(input->rankOf() == 5, 0, "CONV3D CUDNN OP: rank of input array must be equal to 5, but got %i instead !",
input->rankOf());
REQUIRE_TRUE(weights->rankOf() == 5, 0,
"CONV3D CUDNN OP: rank of weights array must be equal to 5, but got %i instead !", weights->rankOf());
LongType kD = INT_ARG(0) > 0 ? INT_ARG(0) : static_cast<LongType>(weights->sizeAt(0)); // filter(kernel) depth
LongType kH = INT_ARG(1) > 0 ? INT_ARG(1) : static_cast<LongType>(weights->sizeAt(1)); // filter(kernel) height
LongType kW = INT_ARG(2) > 0 ? INT_ARG(2) : static_cast<LongType>(weights->sizeAt(2)); // filter(kernel) width
LongType sD = INT_ARG(3); // strides depth
LongType sH = INT_ARG(4); // strides height
LongType sW = INT_ARG(5); // strides width
LongType pD = INT_ARG(6); // paddings depth
LongType pH = INT_ARG(7); // paddings height
LongType pW = INT_ARG(8); // paddings width
LongType dD = INT_ARG(9); // dilations depth
LongType dH = INT_ARG(10); // dilations height
LongType dW = INT_ARG(11); // dilations width
int paddingMode = INT_ARG(12); // 0-SAME, 1-VALID
int isNCDHW = block.getIArguments()->size() > 13 ? !INT_ARG(13) : 1; // INT_ARG(13): 1-NDHWC, 0-NCDHW
int wFormat = block.getIArguments()->size() > 14
? INT_ARG(14)
: 0; // 0-[kD, kH, kW, iC, oC], 1-[oC, iC, kD, kH, kW], 2-[oC, kD, kH, kW, iC]
REQUIRE_TRUE(paddingMode < 2, 0,
"CONV3D CUDNN OP: causal padding mode (paddingMode = 2) is not allowed for this operation !");
LongType bS, iC, iD, iH, iW, oC, oD, oH,
oW; // batch size, input channels, input depth/height/width, output channels, output depth/height/width;
LongType indIOioC, indIOioD, indWoC, indWiC, indWkD; // corresponding indexes
ConvolutionUtils::getSizesAndIndexesConv3d(isNCDHW, wFormat, *input, *output, bS, iC, iD, iH, iW, oC, oD, oH, oW,
indIOioC, indIOioD, indWiC, indWoC, indWkD);
ConvolutionUtils::calcPadding3D(pD, pH, pW, oD, oH, oW, iD, iH, iW, kD, kH, kW, sD, sH, sW, dD, dH, dW, paddingMode);
std::vector<LongType> expectedWeightsShape = ConvolutionUtils::expectWeightsShape(wFormat, kD, kH, kW, iC, oC);
REQUIRE_TRUE(weights->isSameShape(expectedWeightsShape), 0,
"CONV3D CUDNN OP: wrong shape of weights array, expected is %s, but got %s instead !",
ShapeUtils::shapeAsString(expectedWeightsShape).c_str(), ShapeUtils::shapeAsString(weights).c_str());
if (bias)
REQUIRE_TRUE(
bias->rankOf() <= 2 && oC == bias->lengthOf(), 0,
"CONV3D CUDNN OP: wrong shape of array with biases, expected rank, length: <=2, %i, but got %i, %i instead !",
oC, bias->rankOf(), bias->lengthOf());
std::unique_ptr<NDArray> tmpWeight = {}, tmpInput = {};
NDArray* newWeights = weights; // cudnn support only one format {oC,iC,kD,kH,kW}
if (1 != wFormat) {
// Create the tmpWeight object - this syntax is already valid
std::vector<LongType> weightShape = {oC, iC, kD, kH, kW};
// Use the vector for the NDArray constructor
tmpWeight.reset(new NDArray(weights->ordering(), weightShape, weights->dataType(), weights->getContext()));
newWeights = tmpWeight.get();
// Create named vectors as lvalues for the permute call
std::vector<LongType> format0Permute = {4, 3, 0, 1, 2};
std::vector<LongType> format1Permute = {0, 4, 1, 2, 3};
NDArray assign = weights->permute(
0 == wFormat ? format0Permute : format1Permute,
true, // copyToNewBuff
true);
// Use the appropriate one in the permute call
newWeights->assign(&assign); // resetStrides
}
if (paddingMode == 1) { // in same paddingMode cudnn doesn't support asymmetric left/right top/bottopm paddings
auto ret = checkConv3dCUDNNPadAsymmetric(input, nullptr, iD, iH, iW, oD, oH, oW, kD, kH, kW, sD, sH, sW, pD, pH, pW,
dD, dH, dW, isNCDHW);
tmpInput = std::move(std::get<0>(ret)); // prolong life
if (tmpInput) input = tmpInput.get();
}
conv3dCUDNN(block.launchContext(), input, newWeights, bias, output, kD, kH, kW, sD, sH, sW, pD, pH, pW, dD, dH, dW,
paddingMode, isNCDHW, wFormat);
return Status::OK;
}
//////////////////////////////////////////////////////////////////////////
PLATFORM_CHECK(conv3dnew, ENGINE_CUDA) {
auto input = INPUT_VARIABLE(0); // [bS, iD, iH, iW, iC] (NDHWC) or [bS, iC, iD, iH, iW] (NCDHW)
auto weights = INPUT_VARIABLE(1); // [kD, kH, kW, iC, oC], [oC, iC, kD, kH, kW], [oC, kD, kH, kW, iC]
auto bias = block.width() > 2 ? INPUT_VARIABLE(2) : nullptr; // [oC]
int paddingMode = INT_ARG(12); // 0-SAME, 1-VALID
Requirements req("CUDNN CONV3d OP");
req.expectNotEq(makeInfoVariable(paddingMode, "paddingMode"), 2) &&
req.expectIn(makeInfoVariable(input->dataType(), TYPE_MSG_INPUT0),
{HALF, FLOAT32, DOUBLE}) &&
req.expectIn(makeInfoVariable(weights->dataType(), TYPE_MSG_INPUT1),
{HALF, FLOAT32, DOUBLE});
if (bias) {
req.expectIn(makeInfoVariable(bias->dataType(), TYPE_MSG_INPUT_ "#bias"),
{HALF, FLOAT32, DOUBLE});
}
req.logTheSuccess();
return req;
}
//////////////////////////////////////////////////////////////////////////
PLATFORM_IMPL(conv3dnew_bp, ENGINE_CUDA) {
auto input = INPUT_VARIABLE(0); // [bS, iD, iH, iW, iC] (NDHWC) or [bS, iC, iD, iH, iW] (NCDHW)
auto weights = INPUT_VARIABLE(1); // [kD, kH, kW, iC, oC], [oC, iC, kD, kH, kW], [oC, kD, kH, kW, iC]
auto bias = block.width() > 3 ? INPUT_VARIABLE(2) : nullptr; // [oC]
auto gradO = block.width() > 3
? INPUT_VARIABLE(3)
: INPUT_VARIABLE(2); // [bS, oD, oH, oW, oC] (NDHWC) or [bS, oC, oD, oH, oW] (NCDHW), epsilon_next
auto gradI = OUTPUT_VARIABLE(0); // [bS, iD, iH, iW, iC] (NDHWC) or [bS, iC, iD, iH, iW] (NCDHW), epsilon
auto gradW = OUTPUT_VARIABLE(1); // [kD, kH, kW, iC, oC], [oC, iC, kD, kH, kW], [oC, kD, kH, kW, iC]
auto gradB = block.width() > 3 ? OUTPUT_VARIABLE(2) : nullptr; // [oC]
REQUIRE_TRUE(input->rankOf() == 5, 0,
"CONV3D_BP CUDNN OP: rank of input array must be equal to 5, but got %i instead !", input->rankOf());
REQUIRE_TRUE(weights->rankOf() == 5, 0,
"CONV3D_BP CUDNN OP: rank of weights array must be equal to 5, but got %i instead !", weights->rankOf());
REQUIRE_TRUE(
gradO->rankOf() == 5, 0,
"CONV3D_BP CUDNN OP: rank of output gradients (next epsilon) array must be equal to 5, but got %i instead !",
gradO->rankOf());
LongType kD = INT_ARG(0) > 0 ? INT_ARG(0) : static_cast<LongType>(weights->sizeAt(0)); // filter(kernel) depth
LongType kH = INT_ARG(1) > 0 ? INT_ARG(1) : static_cast<LongType>(weights->sizeAt(1)); // filter(kernel) height
LongType kW = INT_ARG(2) > 0 ? INT_ARG(2) : static_cast<LongType>(weights->sizeAt(2)); // filter(kernel) width
LongType sD = INT_ARG(3); // strides depth
LongType sH = INT_ARG(4); // strides height
LongType sW = INT_ARG(5); // strides width
LongType pD = INT_ARG(6); // paddings depth
LongType pH = INT_ARG(7); // paddings height
LongType pW = INT_ARG(8); // paddings width
LongType dD = INT_ARG(9); // dilations depth
LongType dH = INT_ARG(10); // dilations height
LongType dW = INT_ARG(11); // dilations width
int paddingMode = INT_ARG(12); // 1-SAME, 0-VALID
int isNCDHW = block.getIArguments()->size() > 13 ? !INT_ARG(13) : 1; // INT_ARG(13): 1-NDHWC, 0-NCDHW
int wFormat = block.getIArguments()->size() > 14
? INT_ARG(14)
: 0; // 0-[kD, kH, kW, iC, oC], 1-[oC, iC, kD, kH, kW], 2-[oC, kD, kH, kW, iC]
LongType bS, iC, iD, iH, iW, oC, oD, oH,
oW; // batch size, input channels, input depth/height/width, output channels, output depth/height/width;
LongType indIOioC, indIOioD, indWoC, indWiC, indWkD; // corresponding indexes
ConvolutionUtils::getSizesAndIndexesConv3d(isNCDHW, wFormat, *input, *gradO, bS, iC, iD, iH, iW, oC, oD, oH, oW,
indIOioC, indIOioD, indWiC, indWoC, indWkD);
LongType trueoD, trueoH, trueoW; // true output depth/height/width
ConvolutionUtils::calcOutSizePool3D(trueoD, trueoH, trueoW, kD, kH, kW, sD, sH, sW, pD, pH, pW, dD, dH, dW, iD, iH,
iW, paddingMode);
REQUIRE_TRUE(paddingMode < 2, 0,
"CONV3D_BP CUDNN OP: causal padding mode (paddingMode = 2) is not allowed for this operation !");
std::vector<LongType> expectedGradOShape = ShapeUtils::composeShapeUsingDimsAndIdx(
{bS, oC, trueoD, trueoH, trueoW, 0, indIOioC, indIOioD, indIOioD + 1, indIOioD + 2});
std::vector<LongType> expectedWeightsShape = ConvolutionUtils::expectWeightsShape(wFormat, kD, kH, kW, iC, oC);
REQUIRE_TRUE(
gradO->isSameShape(expectedGradOShape), 0,
"CONV3D_BP CUDNN OP: wrong shape of output gradients (next epsilon) array, expected is %s, but got %s instead !",
ShapeUtils::shapeAsString(expectedGradOShape).c_str(), ShapeUtils::shapeAsString(gradO).c_str());
REQUIRE_TRUE(gradW->isSameShape(expectedWeightsShape), 0,
"CONV3D_BP CUDNN OP: wrong shape of weights array, expected is %s, but got %s instead !",
ShapeUtils::shapeAsString(expectedWeightsShape).c_str(), ShapeUtils::shapeAsString(weights).c_str());
if (bias)
REQUIRE_TRUE(bias->rankOf() <= 2 && oC == bias->lengthOf(), 0,
"CONV3D_BP CUDNN OP: wrong shape of array with biases, expected rank, length: <=2, %i, but got %i, %i "
"instead !",
oC, bias->rankOf(), bias->lengthOf());
ConvolutionUtils::calcPadding3D(pD, pH, pW, oD, oH, oW, iD, iH, iW, kD, kH, kW, sD, sH, sW, dD, dH, dW, paddingMode);
std::unique_ptr<NDArray> tmpGradI = {}, tmpInput = {}, tmpWeights = {}, tmpGradW = {};
NDArray *newWeights = weights,
*newGradW = gradW; // cudnn support only two formats {oC,iC,kD,kH,kW} and {oC,kD,kH,kW,iC}
if (0 == wFormat) {
// Create named vectors as lvalues for the NDArray constructor
std::vector<LongType> ncdhwGradShape = {oC, iC, kD, kH, kW};
std::vector<LongType> ndhwcGradShape = {oC, kD, kH, kW, iC};
// Use the appropriate one for the gradW reset
tmpGradW.reset(new NDArray(
gradW->ordering(),
isNCDHW ? ncdhwGradShape : ndhwcGradShape,
gradW->dataType(), gradW->getContext()));
// Create named vectors as lvalues for the NDArray constructor
std::vector<LongType> ncdhwShape = {oC, iC, kD, kH, kW};
std::vector<LongType> ndhwcShape = {oC, kD, kH, kW, iC};
// Use the appropriate one for the weights reset
tmpWeights.reset(new NDArray(
weights->ordering(),
isNCDHW ? ncdhwShape : ndhwcShape,
weights->dataType(), weights->getContext()));
// Create named vectors as lvalues for the permute call
std::vector<LongType> ncdhwPermute = {4, 3, 0, 1, 2};
std::vector<LongType> ndhwcPermute = {4, 0, 1, 2, 3};
// Set the pointer variables
newGradW = tmpGradW.get();
newWeights = tmpWeights.get();
NDArray assign = weights->permute(
isNCDHW ? ncdhwPermute : ndhwcPermute,
true, // copyToNewBuff
true);
// Use the appropriate one in the permute call
newWeights->assign(&assign); // resetStrides (kD, kH, kW, iC, oC --> oC, iC, kD, kH, kW) or (kD, kH, kW,
// iC, oC --> oC, kD, kH, kW, iC)
}
NDArray* newInput = input;
NDArray* newGradI = gradI;
if (paddingMode == 1) { // in same paddingMode cudnn doesn't support asymmetric left/right top/bottopm paddings
auto ret = checkConv3dCUDNNPadAsymmetric(input, gradI, iD, iH, iW, oD, oH, oW, kD, kH, kW, sD, sH, sW, pD, pH, pW,
dD, dH, dW, isNCDHW);
tmpInput = std::move(std::get<0>(ret));
tmpGradI = std::move(std::get<1>(ret));
if (tmpInput) newInput = tmpInput.get();
if (tmpGradI) newGradI = tmpGradI.get();
}
conv3dBpCUDNN(block.launchContext(), newInput, newWeights, gradO, newGradI, newGradW, gradB, kD, kH, kW, sD, sH, sW,
pD, pH, pW, dD, dH, dW, paddingMode, isNCDHW, wFormat);
if (0 == wFormat) {
// Create named vectors as lvalues for the permutei call
std::vector<LongType> ncdhwPermutei = {2, 3, 4, 1, 0};
std::vector<LongType> ndhwcPermutei = {1, 2, 3, 4, 0};
// Use the appropriate one in the permutei call
newGradW->permutei(
isNCDHW ? ncdhwPermutei : ndhwcPermutei,false,false); // (oC, iC, kD, kH, kW --> kD, kH, kW, iC, oC) or (oC, kD, kH, kW, iC --> kD, kH, kW, iC, oC) gradW->assign(newGradW);
}
if (newInput != input) {
if (isNCDHW) {
NDArray assign = (*newGradI)({0, 0, 0, 0, 0, gradI->sizeAt(2), 0, gradI->sizeAt(3), 0, gradI->sizeAt(4)});
gradI->assign(&assign);
} else {
NDArray assign = (*newGradI)({0, 0, 0, gradI->sizeAt(1), 0, gradI->sizeAt(2), 0, gradI->sizeAt(3), 0, 0});
gradI->assign(&assign);
}
}
return Status::OK;
}
PLATFORM_CHECK(conv3dnew_bp, ENGINE_CUDA) {
auto input = INPUT_VARIABLE(0); // [bS, iD, iH, iW, iC] (NDHWC) or [bS, iC, iD, iH, iW] (NCDHW)
auto weights = INPUT_VARIABLE(1); // [kD, kH, kW, iC, oC], [oC, iC, kD, kH, kW], [oC, kD, kH, kW, iC]
auto bias = block.width() > 3 ? INPUT_VARIABLE(2) : nullptr; // [oC]
auto gradO = block.width() > 3
? INPUT_VARIABLE(3)
: INPUT_VARIABLE(2); // [bS, oD, oH, oW, oC] (NDHWC) or [bS, oC, oD, oH, oW] (NCDHW), epsilon_next
LongType paddingMode = INT_ARG(12); // 1-SAME, 0-VALID
int isNCDHW = block.getIArguments()->size() > 13 ? !INT_ARG(13) : 1; // INT_ARG(13): 1-NDHWC, 0-NCDHW
Requirements req("CUDNN CONV3d_BP OP");
req.expectNotEq(makeInfoVariable(paddingMode, "paddingMode"), 2) &&
req.expectTrue(makeInfoVariable(isNCDHW, "isNCDHW")) &&
req.expectIn(makeInfoVariable(input->dataType(), TYPE_MSG_INPUT0),
{HALF, FLOAT32, DOUBLE}) &&
req.expectIn(makeInfoVariable(weights->dataType(), TYPE_MSG_INPUT1),
{HALF, FLOAT32, DOUBLE});
if (bias) {
req.expectIn(makeInfoVariable(bias->dataType(), TYPE_MSG_INPUT_ "#bias"),
{HALF, FLOAT32, DOUBLE}) &&
req.expectIn(makeInfoVariable(gradO->dataType(), TYPE_MSG_INPUT3),
{HALF, FLOAT32, DOUBLE});
} else {
req.expectIn(makeInfoVariable(gradO->dataType(), TYPE_MSG_INPUT2),
{HALF, FLOAT32, DOUBLE});
}
req.logTheSuccess();
return req;
}
} // namespace platforms
} // namespace ops
} // namespace sd
@@ -0,0 +1,189 @@
/*******************************************************************************
*
* Copyright (c) 2021 Konduit K.K.
*
* 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.
*
* 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 AbdelRauf
//
#include <array/NDArrayFactory.h>
#include <vector>
#include "cudnnUtils.h"
namespace sd {
namespace ops {
namespace platforms {
std::vector<int> getConcatTargets(NDArray&targetLabels, NDArray&targetLabelLengths) {
// concatenate target labels
const int32_t *tlabels = bufferInHost<int32_t>(targetLabels);
const int32_t *tlens = bufferInHost<int32_t>(targetLabelLengths);
int32_t nextOffset = targetLabels.strideAt(0);
int32_t elStride = targetLabels.strideAt(1);
int32_t batchCount = targetLabelLengths.lengthOf();
std::vector<int> labels;
labels.resize(targetLabels.lengthOf());
int j = 0;
for (int i = 0; i < batchCount; i++) {
int count = tlens[i];
for (int k = 0; k < count; k++) {
labels[j] = tlabels[k * elStride];
j++;
}
tlabels += nextOffset;
}
return labels;
}
void cudnnCtcLoss(const LaunchContext &context, NDArray&probs, const int32_t *targetLabelsPtr,
NDArray&probInputLengthes, NDArray&targetLabelLengths, NDArray &ctcLosses,
NDArray &grads) {
const int dims[] = {(int)probs.sizeAt(0), (int)probs.sizeAt(1), (int)probs.sizeAt(2)};
const int strides[] = {(int)probs.strideAt(0), (int)probs.strideAt(1), (int)probs.strideAt(2)};
auto handle = reinterpret_cast<cudnnHandle_t *>(context.getCuDnnHandle());
CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnSetStream), cudnnSetStream(*handle, *context.getCudaStream()));
CTCLossDesc ctcLossDesc;
CudnnTensor probsDesc, gradsDesc(nullptr);
bool calcGrads = !grads.isEmpty();
auto cudnnType = cudnnDataType(probs.dataType());
ctcLossDesc.set(cudnnType, CUDNN_LOSS_NORMALIZATION_SOFTMAX, CUDNN_PROPAGATE_NAN);
probsDesc.set(cudnnType, probs.rankOf(), dims, strides);
if (calcGrads) {
gradsDesc.create();
const int gradStrides[] = {(int)grads.strideAt(0), (int)grads.strideAt(1), (int)grads.strideAt(2)};
gradsDesc.set(cudnnDataType(grads.dataType()), grads.rankOf(), dims, gradStrides);
}
size_t tempWorkSpaceSize = 0;
CHECK_CUDNN_FAILURE_MSG(
STRINGIZE(cudnnGetCTCLossWorkspaceSize),
cudnnGetCTCLossWorkspaceSize(*handle, probsDesc, gradsDesc, targetLabelsPtr,
bufferInHost<int32_t>(targetLabelLengths), bufferInHost<int32_t>(probInputLengthes),
CUDNN_CTC_LOSS_ALGO_DETERMINISTIC, ctcLossDesc, &tempWorkSpaceSize));
PointersManager manager(&context, __func__);
// Allocate temp tempWorkspace buffer
void *tempWorkSpace = manager.allocateDevMem(tempWorkSpaceSize);
NDArray::prepareSpecialUse({&ctcLosses, &grads}, {&probs});
CHECK_CUDNN_FAILURE_MSG(
STRINGIZE(cudnnCTCLoss),
cudnnCTCLoss(*handle, probsDesc, probs.specialBuffer(), targetLabelsPtr,
bufferInHost<int32_t>(targetLabelLengths), bufferInHost<int32_t>(probInputLengthes),
ctcLosses.specialBuffer(), gradsDesc, calcGrads ? grads.specialBuffer() : nullptr,
CUDNN_CTC_LOSS_ALGO_DETERMINISTIC, ctcLossDesc, tempWorkSpace, tempWorkSpaceSize));
NDArray::registerSpecialUse({&ctcLosses, &grads}, {&probs});
return;
}
PLATFORM_IMPL(ctc_loss, ENGINE_CUDA) {
auto targetLabels = INPUT_VARIABLE(0);
auto logitInput = INPUT_VARIABLE(1);
auto targetLabelLengths = INPUT_VARIABLE(2);
auto logitInputLengths = INPUT_VARIABLE(3);
auto outputLosses = OUTPUT_VARIABLE(0);
auto context = block.launchContext();
// in Cudnn Batch is in the middle dimension
logitInput->permutei({1, 0, 2});
// in Cudnn targets are concantenated instead of batched as matrix
auto labels = getConcatTargets(*targetLabels, *targetLabelLengths);
const int32_t *ldata = labels.data();
auto emptyGrads = NDArrayFactory::empty<float>();
cudnnCtcLoss(*context, *logitInput, ldata, *logitInputLengths, *targetLabelLengths, *outputLosses, emptyGrads);
return Status::OK;
}
template <typename T>
bool checkLabelLength(NDArray&labelLengthArr) {
// check label lengths
auto lenBatch = labelLengthArr.lengthOf();
for (int i = 0; i < lenBatch; i++) {
// The labelLengths is greater than 256.
if (labelLengthArr.e<int32_t>(i) > 256) return false;
}
return true;
}
PLATFORM_CHECK(ctc_loss, ENGINE_CUDA) {
auto targetLabels = INPUT_VARIABLE(0);
auto logitInput = INPUT_VARIABLE(1);
auto targetLabelLengths = INPUT_VARIABLE(2);
auto logitInputLengths = INPUT_VARIABLE(3);
auto outputLosses = OUTPUT_VARIABLE(0);
int blankIndex = INT_ARG(0);
Requirements req("CUDNN CTC_LOSS OP");
req.expectEq(makeInfoVariable(blankIndex, "Blank Index"), 0) &&
req.expectEq(makeInfoVariable(logitInput->dataType(), TYPE_MSG_INPUT1), FLOAT32) &&
req.expectEq(makeInfoVariable(targetLabelLengths->dataType(), TYPE_MSG_INPUT2), INT32) &&
req.expectTrue(
makeInfoVariable(checkLabelLength<int32_t>(*targetLabelLengths), "target Label lengthes should be <= 256"),
NO_MSG);
req.logTheSuccess();
return req;
}
PLATFORM_IMPL(ctc_loss_grad, ENGINE_CUDA) {
auto targetLabels = INPUT_VARIABLE(0);
auto logitInput = INPUT_VARIABLE(1);
auto targetLabelLengths = INPUT_VARIABLE(2);
auto logitInputLengths = INPUT_VARIABLE(3);
auto outputGradients = OUTPUT_VARIABLE(0);
auto context = block.launchContext();
REQUIRE_TRUE(outputGradients->isSameShape(logitInput), 0,
"CtcLoss Gradient: wrong shape of output array, expected is %s but got %s instead !",
ShapeUtils::shapeAsString(logitInput).c_str(), ShapeUtils::shapeAsString(outputGradients).c_str());
// in Cudnn Batch is in the middle dimension
logitInput->permutei({1, 0, 2});
outputGradients->permutei({1, 0, 2});
// in Cudnn targets are concantenated instead of batched as matrix
auto labels = getConcatTargets(*targetLabels, *targetLabelLengths);
const int32_t *ldata = labels.data();
auto tempLosses = NDArrayFactory::create<float>('c', {logitInputLengths->sizeAt(0)});
cudnnCtcLoss(*context, *logitInput, ldata, *logitInputLengths, *targetLabelLengths, tempLosses, *outputGradients);
// restore grads shape from {T, BATCH, C} -> {BATCHS, T, C}
outputGradients->permutei({1, 0, 2});
return Status::OK;
}
PLATFORM_CHECK(ctc_loss_grad, ENGINE_CUDA) {
auto targetLabels = INPUT_VARIABLE(0);
auto logitInput = INPUT_VARIABLE(1);
auto targetLabelLengths = INPUT_VARIABLE(2);
auto logitInputLengths = INPUT_VARIABLE(3);
auto outputGrads = OUTPUT_VARIABLE(0);
int blankIndex = INT_ARG(0);
Requirements req("CUDNN CTC_LOSS_GRAD OP");
req.expectEq(makeInfoVariable(blankIndex, "Blank Index"), 0) &&
req.expectEq(makeInfoVariable(logitInput->dataType(), TYPE_MSG_INPUT1), FLOAT32) &&
req.expectEq(makeInfoVariable(targetLabelLengths->dataType(), TYPE_MSG_INPUT2), INT32) &&
req.expectTrue(
makeInfoVariable(checkLabelLength<int32_t>(*targetLabelLengths), "target Label lengthes should be <= 256"),
NO_MSG);
req.logTheSuccess();
return req;
}
} // namespace platforms
} // namespace ops
} // namespace sd
@@ -0,0 +1,397 @@
/* ******************************************************************************
*
*
* 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)
//
#include <ops/declarable/helpers/convolutions.h>
#include "cudnnUtils.h"
namespace sd {
namespace ops {
namespace platforms {
//////////////////////////////////////////////////////////////////////////
std::tuple<std::unique_ptr<NDArray>, std::unique_ptr<NDArray>> checkConv2dCUDNNPadAsymmetric(
NDArray* input, NDArray* gradI, const int iH, const int iW, const int oH, const int oW, const int kH,
const int kW, const int sH, const int sW, const int pH, const int pW, const int dH, const int dW,
const bool isNCHW) {
const auto pHsum = ((oH - 1) * sH + ((kH - 1) * dH + 1) - iH);
const auto pWsum = ((oW - 1) * sW + ((kW - 1) * dW + 1) - iW);
const bool isPHasymm = pH != (pHsum - pH);
const bool isPWasymm = pW != (pWsum - pW);
std::unique_ptr<NDArray> uNewInput = {}, uNewGradI = {};
if (!isPHasymm && !isPWasymm) return std::make_tuple(std::move(uNewInput), std::move(uNewGradI));
std::vector<LongType> newShape = input->getShapeAsVector();
const int iHposition = isNCHW ? 2 : 1;
if (isPHasymm) newShape[iHposition] += 1;
if (isPWasymm) newShape[iHposition + 1] += 1;
uNewInput.reset(new NDArray(input->ordering(), newShape, input->dataType(), input->getContext()));
if (isNCHW)
(*uNewInput)({0, 0, 0, 0, 0, input->sizeAt(2), 0, input->sizeAt(3)}).assign(input);
else
(*uNewInput)({0, 0, 0, input->sizeAt(1), 0, input->sizeAt(2), 0, 0}).assign(input);
if (gradI != nullptr)
uNewGradI.reset(new NDArray(gradI->ordering(), newShape, gradI->dataType(), gradI->getContext()));
return std::make_tuple(std::move(uNewInput), std::move(uNewGradI));
}
//////////////////////////////////////////////////////////////////////////
std::tuple<std::unique_ptr<NDArray>, std::unique_ptr<NDArray>> checkConv3dCUDNNPadAsymmetric(
NDArray* input, NDArray* gradI, const int iD, const int iH, const int iW, const int oD, const int oH,
const int oW, const int kD, const int kH, const int kW, const int sD, const int sH, const int sW, const int pD,
const int pH, const int pW, const int dD, const int dH, const int dW, const bool isNCDHW) {
const auto pDsum = ((oD - 1) * sD + ((kD - 1) * dD + 1) - iD);
const auto pHsum = ((oH - 1) * sH + ((kH - 1) * dH + 1) - iH);
const auto pWsum = ((oW - 1) * sW + ((kW - 1) * dW + 1) - iW);
const bool isPDasymm = pD != (pDsum - pD);
const bool isPHasymm = pH != (pHsum - pH);
const bool isPWasymm = pW != (pWsum - pW);
std::unique_ptr<NDArray> uNewInput = {}, uNewGradI = {};
if (!isPDasymm && !isPHasymm && !isPWasymm) return std::make_tuple(std::move(uNewInput), std::move(uNewGradI));
std::vector<LongType> newShape = input->getShapeAsVector();
const int iDposition = isNCDHW ? 2 : 1;
if (isPDasymm) newShape[iDposition] += 1;
if (isPHasymm) newShape[iDposition + 1] += 1;
if (isPWasymm) newShape[iDposition + 2] += 1;
uNewInput.reset(new NDArray(input->ordering(), newShape, input->dataType(), input->getContext()));
if (isNCDHW)
(*uNewInput)({0, 0, 0, 0, 0, input->sizeAt(2), 0, input->sizeAt(3), 0, input->sizeAt(4)}).assign(input);
else
(*uNewInput)({0, 0, 0, input->sizeAt(1), 0, input->sizeAt(2), 0, input->sizeAt(3), 0, 0}).assign(input);
if (gradI != nullptr)
uNewGradI.reset(new NDArray(gradI->ordering(), newShape, gradI->dataType(), gradI->getContext()));
return std::make_tuple(std::move(uNewInput), std::move(uNewGradI));
}
//////////////////////////////////////////////////////////////////////////
void pooling2dCUDNN(const LaunchContext* context, NDArray* input, NDArray* output, const int kH, const int kW,
const int sH, const int sW, const int pH, const int pW, const int dH, const int dW,
const bool isNCHW, const cudnnPoolingMode_t mode) {
LongType bS, iC, iH, iW, oC, oH,
oW; // batch size, input channels, input height/width, output channels, output height/width;
LongType indIOioC, indIiH, indWoC, indWiC, indWkH, indOoH; // corresponding indexes
ConvolutionUtils::getSizesAndIndexesConv2d(isNCHW, 0, *input, *output, bS, iC, iH, iW, oC, oH, oW, indIOioC, indIiH,
indWiC, indWoC, indWkH, indOoH);
auto handle = reinterpret_cast<cudnnHandle_t*>(context->getCuDnnHandle());
CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnSetStream), cudnnSetStream(*handle, *context->getCudaStream()));
cudnnTensorFormat_t format = isNCHW ? CUDNN_TENSOR_NCHW : CUDNN_TENSOR_NHWC;
// input descriptor, output descriptor
CudnnTensor x, z;
if (input->ordering() == 'c')
x.set4D(format, cudnnDataType(input->dataType()), bS, iC, iH, iW);
else
x.set4DEx(cudnnDataType(input->dataType()), bS, iC, iH, iW, input->strideAt(0), input->strideAt(indIOioC),
input->strideAt(indIiH), input->strideAt(indIiH + 1));
if (output->ordering() == 'c')
z.set4D(format, cudnnDataType(output->dataType()), bS, oC, oH, oW);
else
z.set4DEx(cudnnDataType(output->dataType()), bS, oC, oH, oW, output->strideAt(0), output->strideAt(indIOioC),
output->strideAt(indOoH), output->strideAt(indOoH + 1));
// description of pooling
PoolingDesc pooling;
pooling.set2D(mode, CUDNN_PROPAGATE_NAN, kH, kW, pH, pW, sH, sW);
// provide scaling parameters
const float alpha32(1), beta32(0);
const double alpha64(1), beta64(0);
const void* alpha =
output->sizeOfT() <= 4 ? reinterpret_cast<const void*>(&alpha32) : reinterpret_cast<const void*>(&alpha64);
const void* beta =
output->sizeOfT() <= 4 ? reinterpret_cast<const void*>(&beta32) : reinterpret_cast<const void*>(&beta64);
NDArray::prepareSpecialUse({output}, {input});
// run calculation
CHECK_CUDNN_FAILURE_MSG(
STRINGIZE(cudnnPoolingForward),
cudnnPoolingForward(*handle, pooling, alpha, x, input->specialBuffer(), beta, z, output->specialBuffer()));
auto cudaErr = cudaStreamSynchronize(*context->getCudaStream());
if (cudaErr != 0) throw cuda_exception::build("pooling2dCUDNN: cudaStreamSynchronize failed !", cudaErr);
NDArray::registerSpecialUse({output}, {input});
}
//////////////////////////////////////////////////////////////////////////
void pooling2dBpCUDNN(const LaunchContext* context, NDArray* input, NDArray* gradO, NDArray* gradI,
const int kH, const int kW, const int sH, const int sW, const int pH, const int pW, const int dH,
const int dW, const bool isNCHW, const cudnnPoolingMode_t mode) {
LongType bS, iC, iH, iW, oC, oH,
oW; // batch size, input channels, input height/width, output channels, output height/width;
LongType indIOioC, indIiH, indWoC, indWiC, indWkH, indOoH; // corresponding indexes
ConvolutionUtils::getSizesAndIndexesConv2d(isNCHW, 0, *input, *gradO, bS, iC, iH, iW, oC, oH, oW, indIOioC, indIiH,
indWiC, indWoC, indWkH, indOoH);
auto handle = reinterpret_cast<cudnnHandle_t*>(context->getCuDnnHandle());
CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnSetStream), cudnnSetStream(*handle, *context->getCudaStream()));
cudnnTensorFormat_t format = isNCHW ? CUDNN_TENSOR_NCHW : CUDNN_TENSOR_NHWC;
// input and gradI descriptor
CudnnTensor x;
if (input->ordering() == 'c')
x.set4D(format, cudnnDataType(input->dataType()), bS, iC, iH, iW);
else
x.set4DEx(cudnnDataType(input->dataType()), bS, iC, iH, iW, input->strideAt(0), input->strideAt(indIOioC),
input->strideAt(indIiH), input->strideAt(indIiH + 1));
// gradO descriptor
CudnnTensor dz;
if (gradO->ordering() == 'c')
dz.set4D(format, cudnnDataType(gradO->dataType()), bS, oC, oH, oW);
else
dz.set4DEx(cudnnDataType(gradO->dataType()), bS, oC, oH, oW, gradO->strideAt(0), gradO->strideAt(indIOioC),
gradO->strideAt(indOoH), gradO->strideAt(indOoH + 1));
// description of pooling
PoolingDesc pooling;
pooling.set2D(mode, CUDNN_PROPAGATE_NAN, kH, kW, pH, pW, sH, sW);
// provide scaling parameters
const float alpha32(1), beta32(0);
const double alpha64(1), beta64(0);
const void* alpha =
gradO->sizeOfT() <= 4 ? reinterpret_cast<const void*>(&alpha32) : reinterpret_cast<const void*>(&alpha64);
const void* beta =
gradO->sizeOfT() <= 4 ? reinterpret_cast<const void*>(&beta32) : reinterpret_cast<const void*>(&beta64);
NDArray::prepareSpecialUse({gradI}, {input, gradO});
// run calculation for gradI
CHECK_CUDNN_FAILURE_MSG(
STRINGIZE(cudnnPoolingBackward),
cudnnPoolingBackward(*handle, pooling, alpha, dz, gradO->specialBuffer(), dz, gradO->specialBuffer(), x,
input->specialBuffer(), beta, x, gradI->specialBuffer()));
auto cudaErr = cudaStreamSynchronize(*context->getCudaStream());
if (cudaErr != 0) throw cuda_exception::build("pooling2dBpCUDNN: cudaStreamSynchronize failed !", cudaErr);
NDArray::registerSpecialUse({gradI}, {input, gradO});
}
//////////////////////////////////////////////////////////////////////////
void pooling3dCUDNN(const LaunchContext* context, NDArray* input, NDArray* output, const int kD, const int kH,
const int kW, const int sD, const int sH, const int sW, const int pD, const int pH, const int pW,
const int dD, const int dH, const int dW, const bool isNCDHW, const cudnnPoolingMode_t mode) {
auto handle = reinterpret_cast<cudnnHandle_t*>(context->getCuDnnHandle());
CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnSetStream), cudnnSetStream(*handle, *context->getCudaStream()));
const int numDims = 5;
LongType bS, iC, iD, iH, iW, oC, oD, oH,
oW; // batch size, input channels, input depth/height/width, output channels, output depth/height/width;
LongType indIOioC, indIOioD, indWoC, indWiC, indWkD; // corresponding indexes
ConvolutionUtils::getSizesAndIndexesConv3d(isNCDHW, 0, *input, *output, bS, iC, iD, iH, iW, oC, oD, oH, oW, indIOioC,
indIOioD, indWiC, indWoC, indWkD);
const int pSizes[] = {pD, pH, pW};
const int sSizes[] = {sD, sH, sW};
const int kSizes[] = {kD, kH, kW};
const LongType xShape[] = {bS, iC, iD, iH, iW};
const LongType zShape[] = {bS, oC, oD, oH, oW};
const LongType xStrides[] = {(LongType)input->strideAt(0), (LongType)input->strideAt(1), (LongType)input->strideAt(2),
(LongType)input->strideAt(3), (LongType)input->strideAt(4)};
const LongType zStrides[] = {(LongType)output->strideAt(0), (LongType)output->strideAt(1), (LongType)output->strideAt(2),
(LongType)output->strideAt(3), (LongType)output->strideAt(4)};
cudnnTensorFormat_t format = isNCDHW ? CUDNN_TENSOR_NCHW : CUDNN_TENSOR_NHWC;
// input descriptor, output descriptor
CudnnTensor x, z;
if (input->ordering() == 'c') {
int newShape[5];
for(int i = 0; i < 5; i++) {
newShape[i] = static_cast<int>(xShape[i]);
}
x.setEx(format, cudnnDataType(input->dataType()), numDims, newShape);
} else {
int newShape[5];
int newStrides[5];
for(int i = 0; i < 5; i++) {
newShape[i] = static_cast<int>(xShape[i]);
}
for(int i = 0; i < 5; i++) {
newStrides[i] = static_cast<int>(xStrides[i]);
}
x.set(cudnnDataType(input->dataType()), numDims, newShape, newStrides);
}
if (output->ordering() == 'c') {
int newShape[5];
int newStrides[5];
for(int i = 0; i < 5; i++) {
newShape[i] = static_cast<int>(zShape[i]);
}
for(int i = 0; i < 5; i++) {
newStrides[i] = static_cast<int>(zStrides[i]);
}
z.setEx(format, cudnnDataType(output->dataType()), numDims, newShape);
} else {
int newShape[5];
int newStrides[5];
for(int i = 0; i < 5; i++) {
newShape[i] = static_cast<int>(zShape[i]);
}
for(int i = 0; i < 5; i++) {
newStrides[i] = static_cast<int>(zStrides[i]);
}
z.set(cudnnDataType(output->dataType()), numDims, newShape, newStrides);
}
// description of pooling
PoolingDesc pooling;
pooling.set(mode, CUDNN_PROPAGATE_NAN, numDims - 2, kSizes, pSizes, sSizes);
// provide scaling parameters
const float alpha32(1), beta32(0);
const double alpha64(1), beta64(0);
const void* alpha =
output->sizeOfT() <= 4 ? reinterpret_cast<const void*>(&alpha32) : reinterpret_cast<const void*>(&alpha64);
const void* beta =
output->sizeOfT() <= 4 ? reinterpret_cast<const void*>(&beta32) : reinterpret_cast<const void*>(&beta64);
NDArray::prepareSpecialUse({output}, {input});
// run calculation
CHECK_CUDNN_FAILURE_MSG(
STRINGIZE(cudnnPoolingForward),
cudnnPoolingForward(*handle, pooling, alpha, x, input->specialBuffer(), beta, z, output->specialBuffer()));
auto cudaErr = cudaStreamSynchronize(*context->getCudaStream());
if (cudaErr != 0) throw cuda_exception::build("pooling3dCUDNN: cudaStreamSynchronize failed !", cudaErr);
NDArray::registerSpecialUse({output}, {input});
}
//////////////////////////////////////////////////////////////////////////
void pooling3dBpCUDNN(const LaunchContext* context, NDArray* input, NDArray* gradO, NDArray* gradI,
const int kD, const int kH, const int kW, const int sD, const int sH, const int sW, const int pD,
const int pH, const int pW, const int dD, const int dH, const int dW, const bool isNCDHW,
const cudnnPoolingMode_t mode) {
auto handle = reinterpret_cast<cudnnHandle_t*>(context->getCuDnnHandle());
CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnSetStream), cudnnSetStream(*handle, *context->getCudaStream()));
const int numDims = 5;
LongType bS, iC, iD, iH, iW, oC, oD, oH,
oW; // batch size, input channels, input depth/height/width, output channels, output depth/height/width;
LongType indIOioC, indIOioD, indWoC, indWiC, indWkD; // corresponding indexes
ConvolutionUtils::getSizesAndIndexesConv3d(isNCDHW, 0, *input, *gradO, bS, iC, iD, iH, iW, oC, oD, oH, oW, indIOioC,
indIOioD, indWiC, indWoC, indWkD);
const int pSizes[] = {pD, pH, pW};
const int sSizes[] = {sD, sH, sW};
const int kSizes[] = {kD, kH, kW};
const int xShape[] = {(int) bS, (int) iC, (int) iD, (int) iH, (int) iW};
const int dzShape[] = {(int) bS, (int) oC, (int) oD, (int) oH,(int) oW};
const int xStrides[] = { (int) input->strideAt(0), (int)input->strideAt(1), (int)input->strideAt(2),
(int)input->strideAt(3), (int)input->strideAt(4)};
const int dzStrides[] = {(int)gradO->strideAt(0), (int)gradO->strideAt(1), (int)gradO->strideAt(2),
(int)gradO->strideAt(3), (int)gradO->strideAt(4)};
cudnnTensorFormat_t format = isNCDHW ? CUDNN_TENSOR_NCHW : CUDNN_TENSOR_NHWC;
// input and gradI descriptor
CudnnTensor x;
if ( input->ordering() == 'c') {
x.setEx(format, cudnnDataType(input->dataType()), numDims, xShape);
} else {
x.set(cudnnDataType(input->dataType()), numDims, xShape, xStrides);
}
// gradO descriptor
CudnnTensor dz;
if ( gradO->ordering() == 'c')
dz.setEx(format, cudnnDataType(gradO->dataType()), numDims, dzShape);
else
dz.set(cudnnDataType(gradO->dataType()), numDims, dzShape, dzStrides);
// description of pooling
PoolingDesc pooling;
pooling.set(mode, CUDNN_PROPAGATE_NAN, numDims - 2, kSizes, pSizes, sSizes);
// provide scaling parameters
const float alpha32(1), beta32(0);
const double alpha64(1), beta64(0);
const void* alpha =
gradO->sizeOfT() <= 4 ? reinterpret_cast<const void*>(&alpha32) : reinterpret_cast<const void*>(&alpha64);
const void* beta =
gradO->sizeOfT() <= 4 ? reinterpret_cast<const void*>(&beta32) : reinterpret_cast<const void*>(&beta64);
// cudnn maxpool2d_bp api requires ff output as one of input arguments
if (mode == CUDNN_POOLING_MAX) {
NDArray temp(gradO);
NDArray::prepareSpecialUse({gradI}, {input, gradO, &temp});
// run ff calculation
CHECK_CUDNN_FAILURE_MSG(
STRINGIZE(cudnnPoolingForward),
cudnnPoolingForward(*handle, pooling, alpha, x, input->specialBuffer(), beta, dz, temp.specialBuffer()));
// run bp calculation for gradI
CHECK_CUDNN_FAILURE_MSG(
STRINGIZE(cudnnPoolingBackward),
cudnnPoolingBackward(*handle, pooling, alpha, dz, temp.specialBuffer(), dz, gradO->specialBuffer(), x,
input->specialBuffer(), beta, x, gradI->specialBuffer()));
NDArray::registerSpecialUse({gradI}, {input, gradO, &temp});
} else {
NDArray::prepareSpecialUse({gradI}, {input, gradO});
// run bp calculation for gradI
CHECK_CUDNN_FAILURE_MSG(
STRINGIZE(cudnnPoolingBackward),
cudnnPoolingBackward(*handle, pooling, alpha, dz, gradO->specialBuffer(), dz, gradO->specialBuffer(), x,
input->specialBuffer(), beta, x, gradI->specialBuffer()));
NDArray::registerSpecialUse({gradI}, {input, gradO});
}
auto cudaErr = cudaStreamSynchronize(*context->getCudaStream());
if (cudaErr != 0) throw cuda_exception::build("pooling3dBpCUDNN: cudaStreamSynchronize failed !", cudaErr);
}
} // namespace platforms
} // namespace ops
} // namespace sd
@@ -0,0 +1,358 @@
/* ******************************************************************************
*
*
* This program and the accompanying materials are made available under the
* terms of the Apache License, Version 2.0 which is available at
* https://www.apache.org/licenses/LICENSE-2.0.
*
* See the NOTICE file distributed with this work for additional
* information regarding copyright ownership.
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
* WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
* License for the specific language governing permissions and limitations
* under the License.
*
* SPDX-License-Identifier: Apache-2.0
******************************************************************************/
//
// @author raver119@gmail.com
//
#ifndef SD_CUDNNUTILS_H
#define SD_CUDNNUTILS_H
#include <cudnn.h>
#include <exceptions/cuda_exception.h>
#include <exceptions/datatype_exception.h>
#include <helpers/PointersManager.h>
#include <ops/declarable/OpRegistrator.h>
#include <ops/declarable/PlatformHelper.h>
#include <system/platform_boilerplate.h>
#include <memory>
#include <tuple>
#include <vector>
#define CUDNN_NEW_RNN_API_VER 8001
#define CUDNN_CLIPPING_API_VER 7201
namespace sd {
namespace ops {
namespace platforms {
DECLARE_PLATFORM(conv2d, ENGINE_CUDA);
DECLARE_PLATFORM(conv2d_bp, ENGINE_CUDA);
DECLARE_PLATFORM(conv3dnew, ENGINE_CUDA);
DECLARE_PLATFORM(conv3dnew_bp, ENGINE_CUDA);
DECLARE_PLATFORM(depthwise_conv2d, ENGINE_CUDA);
DECLARE_PLATFORM(depthwise_conv2d_bp, ENGINE_CUDA);
DECLARE_PLATFORM(batchnorm, ENGINE_CUDA);
DECLARE_PLATFORM(batchnorm_bp, ENGINE_CUDA);
DECLARE_PLATFORM(avgpool2d, ENGINE_CUDA);
DECLARE_PLATFORM(avgpool2d_bp, ENGINE_CUDA);
DECLARE_PLATFORM(maxpool2d, ENGINE_CUDA);
DECLARE_PLATFORM(maxpool2d_bp, ENGINE_CUDA);
DECLARE_PLATFORM(avgpool3dnew, ENGINE_CUDA);
DECLARE_PLATFORM(avgpool3dnew_bp, ENGINE_CUDA);
DECLARE_PLATFORM(maxpool3dnew, ENGINE_CUDA);
DECLARE_PLATFORM(maxpool3dnew_bp, ENGINE_CUDA);
DECLARE_PLATFORM(lstmLayer, ENGINE_CUDA);
DECLARE_PLATFORM(ctc_loss, ENGINE_CUDA);
DECLARE_PLATFORM(ctc_loss_grad, ENGINE_CUDA);
//////////////////////////////////////////////////////////////////////////
inline void throwIfCudnnFailed(cudnnStatus_t result_status,
const char* message = "Cudnn error: ", const char* prefix = nullptr) {
if (result_status != CUDNN_STATUS_SUCCESS) {
std::string err_message;
if (prefix) err_message = std::string(prefix) + ": ";
err_message += std::string(message);
throw cuda_exception::build(err_message, result_status);
}
}
#define STRINGIZE(x) STRINGIZE2(x)
#define STRINGIZE2(x) #x
#define CHECK_CUDNN_FAILURE(result_status) throwIfCudnnFailed(result_status, "")
#define CHECK_CUDNN_FAILURE_MSG(custom_message, result_status) \
throwIfCudnnFailed(result_status, custom_message, __func__)
template <typename T>
SD_INLINE const T* bufferInHost(NDArray& array) {
array.syncToHost();
return reinterpret_cast<const T*>(array.buffer());
}
#define MOVEONLY_DESC_IMPL(DESC) \
DESC(const DESC& s) = delete; \
DESC& operator=(const DESC& other) = delete; \
DESC(DESC&& other) noexcept : desc(std::move(other.desc)) { other.desc = {}; } \
DESC& operator=(DESC&& other) noexcept { \
if (&other == this) return *this; \
destroy(); \
desc = std::move(other.desc); \
other.desc = {}; \
return *this; \
}
#define MOVEONLY_DESC_FULL_IMPL(DESC_CLASS, DESC_NAME) \
DESC_CLASS() { create(); } \
DESC_CLASS(cudnn##DESC_NAME##_t created) { desc = created; } \
void create() { CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnCreate##DESC_NAME), cudnnCreate##DESC_NAME(&desc)); } \
void destroy() { \
if (desc) CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnCreate##DESC_NAME), cudnnDestroy##DESC_NAME(desc)); \
desc = {}; \
} \
MOVEONLY_DESC_IMPL(DESC_CLASS) \
operator cudnn##DESC_NAME##_t() const { return desc; } \
~DESC_CLASS() { destroy(); } \
cudnn##DESC_NAME##_t desc;
struct CudnnTensor {
MOVEONLY_DESC_FULL_IMPL(CudnnTensor, TensorDescriptor)
template <typename... Args>
void set(Args&&... args) {
CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnSetTensorNdDescriptor),
cudnnSetTensorNdDescriptor(desc, std::forward<Args>(args)...));
}
template <typename... Args>
void setEx(Args&&... args) {
CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnSetTensorNdDescriptorEx),
cudnnSetTensorNdDescriptorEx(desc, std::forward<Args>(args)...));
}
template <typename... Args>
void set4D(Args&&... args) {
CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnSetTensor4dDescriptor),
cudnnSetTensor4dDescriptor(desc, std::forward<Args>(args)...));
}
template <typename... Args>
void set4DEx(Args&&... args) {
CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnSetTensor4dDescriptorEx),
cudnnSetTensor4dDescriptorEx(desc, std::forward<Args>(args)...));
}
};
struct CudnnTensorList {
MOVEONLY_DESC_IMPL(CudnnTensorList)
CudnnTensorList(int size) {
desc.resize(size);
for (int i = 0; i < size; i++) {
CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnCreateTensorDescriptor), cudnnCreateTensorDescriptor(&desc[i]));
}
}
template <typename... Args>
void set(int index, Args&&... args) {
if (index < desc.size()) {
CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnSetTensorNdDescriptor),
cudnnSetTensorNdDescriptor(desc[index], std::forward<Args>(args)...));
}
}
cudnnTensorDescriptor_t get(int i) const {
if (i < desc.size()) return desc[i];
return nullptr;
}
const cudnnTensorDescriptor_t* getDescriptors() const { return desc.data(); }
void destroy() {
for (auto x : desc) {
CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnDestroyTensorDescriptor), cudnnDestroyTensorDescriptor(x));
}
desc = {};
}
~CudnnTensorList() { destroy(); }
std::vector<cudnnTensorDescriptor_t> desc;
};
struct FilterDesc {
MOVEONLY_DESC_FULL_IMPL(FilterDesc, FilterDescriptor)
template <typename... Args>
void set(Args&&... args) {
CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnSetFilterNdDescriptor),
cudnnSetFilterNdDescriptor(desc, std::forward<Args>(args)...));
}
template <typename... Args>
void set4D(Args&&... args) {
CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnSetFilter4dDescriptor),
cudnnSetFilter4dDescriptor(desc, std::forward<Args>(args)...));
}
};
struct DropoutDesc {
MOVEONLY_DESC_FULL_IMPL(DropoutDesc, DropoutDescriptor)
template <typename... Args>
void set(Args&&... args) {
CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnSetDropoutDescriptor),
cudnnSetDropoutDescriptor(desc, std::forward<Args>(args)...));
}
};
#if CUDNN_VERSION > CUDNN_NEW_RNN_API_VER
struct RnnDataDesc {
MOVEONLY_DESC_FULL_IMPL(RnnDataDesc, RNNDataDescriptor)
template <typename... Args>
void set(Args&&... args) {
CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnSetRNNDataDescriptor),
cudnnSetRNNDataDescriptor(desc, std::forward<Args>(args)...));
}
};
#endif
SD_INLINE void setRnnDescriptorOldApi(cudnnRNNDescriptor_t rnnDesc, cudnnHandle_t handle, cudnnRNNInputMode_t inputMode,
cudnnDirectionMode_t dirMode, cudnnRNNMode_t cellMode, cudnnRNNAlgo_t algo,
cudnnDataType_t mathPrec, int32_t hiddenSize, int32_t numLayers,
cudnnDropoutDescriptor_t dropoutDesc, bool use_tensor_op = false) {
auto err = cudnnSetRNNDescriptor_v6(handle, rnnDesc, hiddenSize, numLayers, dropoutDesc, inputMode, dirMode, cellMode,
algo, mathPrec);
CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnSetRNNDescriptor_v6), err);
#if CUDNN_VERSION >= 7001
if (cudnnGetVersion() >= 7001) {
cudnnMathType_t mathType = use_tensor_op ? CUDNN_TENSOR_OP_MATH : CUDNN_DEFAULT_MATH;
CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnSetRNNMatrixMathType), cudnnSetRNNMatrixMathType(rnnDesc, mathType));
}
#endif
return;
}
struct RnnDesc {
MOVEONLY_DESC_FULL_IMPL(RnnDesc, RNNDescriptor)
template <typename... Args>
void setUsingOldAPI(Args&&... args) {
setRnnDescriptorOldApi(desc, std::forward<Args>(args)...);
}
#if CUDNN_VERSION >= CUDNN_NEW_RNN_API_VER
template <typename... Args>
void set(Args&&... args) {
CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnSetRNNDescriptor_v8),
cudnnSetRNNDescriptor_v8(desc, std::forward<Args>(args)...));
}
#endif
};
struct CTCLossDesc {
MOVEONLY_DESC_FULL_IMPL(CTCLossDesc, CTCLossDescriptor)
template <typename... Args>
void set(Args&&... args) {
CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnSetCTCLossDescriptorEx),
cudnnSetCTCLossDescriptorEx(desc, std::forward<Args>(args)...));
}
};
//////////////////////////////////////////////////////////////////////////
struct PoolingDesc {
MOVEONLY_DESC_FULL_IMPL(PoolingDesc, PoolingDescriptor)
template <typename... Args>
void set(Args&&... args) {
CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnSetPoolingNdDescriptor),
cudnnSetPoolingNdDescriptor(desc, std::forward<Args>(args)...));
}
template <typename... Args>
void set2D(Args&&... args) {
CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnSetPooling2dDescriptor),
cudnnSetPooling2dDescriptor(desc, std::forward<Args>(args)...));
}
};
//////////////////////////////////////////////////////////////////////////
struct ConvolutionDesc {
MOVEONLY_DESC_FULL_IMPL(ConvolutionDesc, ConvolutionDescriptor)
template <typename... Args>
void set(Args&&... args) {
CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnSetConvolutionNdDescriptor),
cudnnSetConvolutionNdDescriptor(desc, std::forward<Args>(args)...));
}
template <typename... Args>
void set2D(Args&&... args) {
CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnSetConvolution2dDescriptor),
cudnnSetConvolution2dDescriptor(desc, std::forward<Args>(args)...));
}
};
//////////////////////////////////////////////////////////////////////////
SD_INLINE cudnnDataType_t cudnnDataType(DataType dataType) {
switch (dataType) {
case FLOAT32:
return CUDNN_DATA_FLOAT;
case DOUBLE:
return CUDNN_DATA_DOUBLE;
case HALF:
return CUDNN_DATA_HALF;
case INT32:
return CUDNN_DATA_INT32;
case INT8:
return CUDNN_DATA_INT8;
default:
throw datatype_exception::build("Unsupported data type", dataType);
}
}
//////////////////////////////////////////////////////////////////////////
std::tuple<std::unique_ptr<NDArray>, std::unique_ptr<NDArray>> checkConv2dCUDNNPadAsymmetric(
NDArray* input, NDArray* gradI, const int iH, const int iW, const int oH, const int oW, const int kH,
const int kW, const int sH, const int sW, const int pH, const int pW, const int dH, const int dW,
const bool isNCHW);
//////////////////////////////////////////////////////////////////////////
std::tuple<std::unique_ptr<NDArray>, std::unique_ptr<NDArray>> checkConv3dCUDNNPadAsymmetric(
NDArray* input, NDArray* gradI, const int iD, const int iH, const int iW, const int oD, const int oH,
const int oW, const int kD, const int kH, const int kW, const int sD, const int sH, const int sW, const int pD,
const int pH, const int pW, const int dD, const int dH, const int dW, const bool isNCDHW);
//////////////////////////////////////////////////////////////////////////
void pooling2dCUDNN(const LaunchContext* context, NDArray* input, NDArray* output, const int kH, const int kW,
const int sH, const int sW, const int pH, const int pW, const int dH, const int dW,
const bool isNCHW, const cudnnPoolingMode_t mode);
//////////////////////////////////////////////////////////////////////////
void pooling2dBpCUDNN(const LaunchContext* context, NDArray* input, NDArray* gradO, NDArray* gradI,
const int kH, const int kW, const int sH, const int sW, const int pH, const int pW, const int dH,
const int dW, const bool isNCHW, const cudnnPoolingMode_t mode);
//////////////////////////////////////////////////////////////////////////
void pooling3dCUDNN(const LaunchContext* context, NDArray* input, NDArray* output, const int kD, const int kH,
const int kW, const int sD, const int sH, const int sW, const int pD, const int pH, const int pW,
const int dD, const int dH, const int dW, const bool isNCDHW, const cudnnPoolingMode_t mode);
//////////////////////////////////////////////////////////////////////////
void pooling3dBpCUDNN(const LaunchContext* context, NDArray* input, NDArray* gradO, NDArray* gradI,
const int kD, const int kH, const int kW, const int sD, const int sH, const int sW, const int pD,
const int pH, const int pW, const int dD, const int dH, const int dW, const bool isNCDHW,
const cudnnPoolingMode_t mode);
} // namespace platforms
} // namespace ops
} // namespace sd
#endif // SD_CUDNNUTILS_H
@@ -0,0 +1,533 @@
/* ******************************************************************************
*
*
* 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)
//
#include <ops/declarable/helpers/convolutions.h>
#include "cudnnUtils.h"
namespace sd {
namespace ops {
namespace platforms {
//////////////////////////////////////////////////////////////////////////
static void depthwiseConv2dCUDNN(const LaunchContext* context, NDArray* input, NDArray* weights,
NDArray* bias, NDArray* output, const LongType kH, const LongType kW, const LongType sH,
const LongType sW, const LongType pH, const LongType pW, const LongType dH, const LongType dW,
const LongType paddingMode, const bool isNCHW) {
// cudnn supports only following case: mC = 1, oC = iC (groupCount == iC)
// input [bS, iC, iH, iW] nchw or [bS, iH, iW, iC] nhwc
// weights [iC, mC, kH, kW]
// bias [oC], may be nullptr
// output [bS, oC, oH, oW] nchw or [bS, oH, oW, oC] nhwc
// oC = iC*mC
LongType bS, iC, iH, iW, mC, oC, oH,
oW; // batch size, input channels, input height/width, output channels, output height/width;
LongType indIOioC, indIiH, indWmC, indWiC, indWkH, indOoH; // corresponding indexes
ConvolutionUtils::getSizesAndIndexesConv2d(isNCHW, 0, *input, *output, bS, iC, iH, iW, oC, oH, oW, indIOioC, indIiH,
indWiC, indWmC, indWkH, indOoH);
mC = weights->sizeAt(1);
auto handle = reinterpret_cast<cudnnHandle_t*>(context->getCuDnnHandle());
CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnSetStream), cudnnSetStream(*handle, *context->getCudaStream()));
cudnnTensorFormat_t format = isNCHW ? CUDNN_TENSOR_NCHW : CUDNN_TENSOR_NHWC;
PointersManager manager(context, __func__);
// input descriptor
CudnnTensor x;
if (input->ordering() == 'c')
x.set4D(format, cudnnDataType(input->dataType()), bS, iC, iH, iW);
else
x.set4DEx(cudnnDataType(input->dataType()), bS, iC, iH, iW, input->strideAt(0), input->strideAt(indIOioC),
input->strideAt(indIiH), input->strideAt(indIiH + 1));
// weights descriptor
FilterDesc w;
w.set4D(cudnnDataType(weights->dataType()), CUDNN_TENSOR_NCHW, iC, mC, kH, kW);
// output descriptor
CudnnTensor z;
if (output->ordering() == 'c')
z.set4D(format, cudnnDataType(output->dataType()), bS, oC, oH, oW);
else
z.set4DEx(cudnnDataType(output->dataType()), bS, oC, oH, oW, output->strideAt(0), output->strideAt(indIOioC),
output->strideAt(indOoH), output->strideAt(indOoH + 1));
// description of convolution
ConvolutionDesc conv;
conv.set2D(pH, pW, sH, sW, dH, dW, CUDNN_CROSS_CORRELATION, cudnnDataType(output->dataType()));
CHECK_CUDNN_FAILURE_MSG(
STRINGIZE(cudnnSetConvolutionGroupCount),
cudnnSetConvolutionGroupCount(
conv, iC)); // set number of groups (depthwise mode) in description of convolution, groupCount == iC
// algorithm description
cudnnConvolutionFwdAlgo_t algo;
cudnnConvolutionFwdAlgoPerf_t algoPerf;
int count = 0;
// CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnGetConvolutionForwardAlgorithm), cudnnGetConvolutionForwardAlgorithm(
// *handle, x, w, conv, z, CUDNN_CONVOLUTION_FWD_PREFER_FASTEST, 0, &algo));
CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnFindConvolutionForwardAlgorithm),
cudnnFindConvolutionForwardAlgorithm(*handle, x, w, conv, z, 1, &count, &algoPerf));
if (count == 0)
throw cuda_exception::build("depthwiseConv2dCUDNN: cudnnGetConvolutionForwardAlgorithm failed", 0);
algo = algoPerf.algo;
// allocate auxiliary device memory, abbreviation ws means workspace
size_t wsSize;
CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnGetConvolutionForwardWorkspaceSize),
cudnnGetConvolutionForwardWorkspaceSize(*handle, x, w, conv, z, algo, &wsSize));
void* wsData = manager.allocateDevMem(wsSize);
// provide scaling parameters
const float alpha32(1), beta32(0);
const double alpha64(1), beta64(0);
const void* alpha =
output->sizeOfT() <= 4 ? reinterpret_cast<const void*>(&alpha32) : reinterpret_cast<const void*>(&alpha64);
const void* beta =
output->sizeOfT() <= 4 ? reinterpret_cast<const void*>(&beta32) : reinterpret_cast<const void*>(&beta64);
NDArray::prepareSpecialUse({output}, {input, weights, bias});
// run calculation
CHECK_CUDNN_FAILURE_MSG(
STRINGIZE(cudnnConvolutionForward),
cudnnConvolutionForward(*handle, alpha, x, input->specialBuffer(), w, weights->specialBuffer(), conv, algo,
wsData, wsSize, beta, z, output->specialBuffer()));
// add bias if it is present
if (bias != nullptr) {
CudnnTensor b;
// b.set( format, cudnnDataType(bias->dataType()), 1, isNCHW ? bias->lengthOf() : 1, 1, isNCHW ? 1:
// bias->lengthOf());
b.set4D(CUDNN_TENSOR_NCHW, cudnnDataType(bias->dataType()), 1, oC, 1, 1);
CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnAddTensor), cudnnAddTensor(*handle, alpha, b, bias->specialBuffer(), alpha,
z, output->specialBuffer()));
}
}
//////////////////////////////////////////////////////////////////////////
static void depthwiseConv2dBpCUDNN(const LaunchContext* context, NDArray* input, NDArray* weights,
NDArray* gradO, NDArray* gradI, NDArray* gradW, NDArray* gradB, const LongType kH,
const LongType kW, const LongType sH, const LongType sW, const LongType pH, const LongType pW, const LongType dH,
const LongType dW, const LongType paddingMode, const bool isNCHW) {
// cudnn supports only following case: mC = 1, oC = iC (groupCount == iC)
// input, gradI [bS, iC, iH, iW] nchw or [bS, iH, iW, iC] nhwc
// weights, gradW [iC, mC, kH, kW]
// gradB [oC], may be nullptr
// gradO [bS, oC, oH, oW] nchw or [bS, oH, oW, oC] nhwc
// oC = iC*mC
LongType bS, iC, iH, iW, mC, oC, oH,
oW; // batch size, input channels, input height/width, output channels, output height/width;
LongType indIOioC, indIiH, indWmC, indWiC, indWkH, indOoH; // corresponding indexes
ConvolutionUtils::getSizesAndIndexesConv2d(isNCHW, 0, *input, *gradO, bS, iC, iH, iW, oC, oH, oW, indIOioC, indIiH,
indWiC, indWmC, indWkH, indOoH);
mC = weights->sizeAt(1);
auto handle = reinterpret_cast<cudnnHandle_t*>(context->getCuDnnHandle());
CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnSetStream), cudnnSetStream(*handle, *context->getCudaStream()));
cudnnTensorFormat_t format = isNCHW ? CUDNN_TENSOR_NCHW : CUDNN_TENSOR_NHWC;
PointersManager manager(context, __func__);
// input descriptor
CudnnTensor x;
if (input->ordering() == 'c')
x.set4D(format, cudnnDataType(input->dataType()), bS, iC, iH, iW);
else
x.set4DEx(cudnnDataType(input->dataType()), bS, iC, iH, iW, input->strideAt(0), input->strideAt(indIOioC),
input->strideAt(indIiH), input->strideAt(indIiH + 1));
// gradO descriptor
CudnnTensor dz;
if (gradO->ordering() == 'c')
dz.set4D(format, cudnnDataType(gradO->dataType()), bS, oC, oH, oW);
else
dz.set4DEx(cudnnDataType(gradO->dataType()), bS, oC, oH, oW, gradO->strideAt(0), gradO->strideAt(indIOioC),
gradO->strideAt(indOoH), gradO->strideAt(indOoH + 1));
// gradI descriptor
CudnnTensor dx;
if (gradI->ordering() == 'c')
dx.set4D(format, cudnnDataType(gradI->dataType()), bS, iC, iH, iW);
else
dx.set4DEx(cudnnDataType(gradI->dataType()), bS, iC, iH, iW, gradI->strideAt(0), gradI->strideAt(indIOioC),
gradI->strideAt(indIiH), gradI->strideAt(indIiH + 1));
// gradW descriptor
FilterDesc dw;
dw.set4D(cudnnDataType(gradW->dataType()), CUDNN_TENSOR_NCHW, iC, mC, kH, kW);
// description of convolution
ConvolutionDesc conv;
conv.set2D(pH, pW, sH, sW, dH, dW, CUDNN_CROSS_CORRELATION, cudnnDataType(gradO->dataType()));
CHECK_CUDNN_FAILURE_MSG(
STRINGIZE(cudnnSetConvolutionGroupCount),
cudnnSetConvolutionGroupCount(
conv, iC)); // set number of groups (depthwise mode) in description of convolution, groupCount == iC
// gradW algorithm description
cudnnConvolutionBwdFilterAlgo_t algoGradW;
cudnnConvolutionBwdFilterAlgoPerf_t algoGradWPerf;
int count = 0;
// CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnGetConvolutionBackwardFilterAlgorithm),
// cudnnGetConvolutionBackwardFilterAlgorithm( *handle, x, dz, conv, dw, CUDNN_CONVOLUTION_BWD_FILTER_PREFER_FASTEST,
// 0, &algoGradW));
CHECK_CUDNN_FAILURE_MSG(
STRINGIZE(cudnnFindConvolutionBackwardFilterAlgorithm),
cudnnFindConvolutionBackwardFilterAlgorithm(*handle, x, dz, conv, dw, 1, &count, &algoGradWPerf));
if (count == 0)
throw cuda_exception::build(
"depthwiseConv2dBpCUDNN: cudnnGetConvolutionBackwardFilterAlgorithm failed as the count is 0 ", 0);
algoGradW = algoGradWPerf.algo;
// gradI algorithm description
cudnnConvolutionBwdDataAlgo_t algoGradI;
cudnnConvolutionBwdDataAlgoPerf_t algoGradIPerf;
CHECK_CUDNN_FAILURE_MSG(
STRINGIZE(cudnnFindConvolutionBackwardDataAlgorithm),
cudnnFindConvolutionBackwardDataAlgorithm(*handle, dw, dz, conv, x, 1, &count, &algoGradIPerf));
if (count == 0)
throw cuda_exception::build(
"depthwiseConv2dBpCUDNN: cudnnGetConvolutionBackwardDataAlgorithm failed as the count is 0 ", 0);
algoGradI = algoGradIPerf.algo;
// allocate auxiliary device memory for gradW calculation, abbreviation ws means workspace
size_t wsGradWSize;
CHECK_CUDNN_FAILURE_MSG(
STRINGIZE(cudnnGetConvolutionBackwardFilterWorkspaceSize),
cudnnGetConvolutionBackwardFilterWorkspaceSize(*handle, x, dz, conv, dw, algoGradW, &wsGradWSize));
void* wsGradWData = manager.allocateDevMem(wsGradWSize);
// allocate auxiliary device memory for gradI calculation, abbreviation ws means workspace
size_t wsGradISize;
CHECK_CUDNN_FAILURE_MSG(
STRINGIZE(cudnnGetConvolutionBackwardDataWorkspaceSize),
cudnnGetConvolutionBackwardDataWorkspaceSize(*handle, dw, dz, conv, dx, algoGradI, &wsGradISize));
void* wsGradIData = manager.allocateDevMem(wsGradISize);
// provide scaling parameters
const float alpha32(1), beta32(0);
const double alpha64(1), beta64(0);
const void* alpha =
gradO->sizeOfT() <= 4 ? reinterpret_cast<const void*>(&alpha32) : reinterpret_cast<const void*>(&alpha64);
const void* beta =
gradO->sizeOfT() <= 4 ? reinterpret_cast<const void*>(&beta32) : reinterpret_cast<const void*>(&beta64);
NDArray::prepareSpecialUse({gradI, gradW, gradB}, {input, weights, gradO});
// run calculation for gradB (if not nullptr)
if (gradB != nullptr) {
CudnnTensor db;
// db.set( format, cudnnDataType(gradB->dataType()), 1, isNCHW ? gradB->lengthOf() : 1, 1, isNCHW ? 1:
// gradB->lengthOf());
db.set4D(CUDNN_TENSOR_NCHW, cudnnDataType(gradB->dataType()), 1, oC, 1, 1);
CHECK_CUDNN_FAILURE_MSG(
STRINGIZE(cudnnConvolutionBackwardBias),
cudnnConvolutionBackwardBias(*handle, alpha, dz, gradO->specialBuffer(), beta, db, gradB->specialBuffer()));
}
// run calculation for gradW
CHECK_CUDNN_FAILURE_MSG(
STRINGIZE(cudnnConvolutionBackwardFilter),
cudnnConvolutionBackwardFilter(*handle, alpha, x, input->specialBuffer(), dz, gradO->specialBuffer(), conv,
algoGradW, wsGradWData, wsGradWSize, beta, dw, gradW->specialBuffer()));
// run calculation for gradI
CHECK_CUDNN_FAILURE_MSG(
STRINGIZE(cudnnConvolutionBackwardData),
cudnnConvolutionBackwardData(*handle, alpha, dw, weights->specialBuffer(), dz, gradO->specialBuffer(), conv,
algoGradI, wsGradIData, wsGradISize, beta, dx, gradI->specialBuffer()));
NDArray::registerSpecialUse({gradI, gradW, gradB}, {input, weights, gradO});
}
//////////////////////////////////////////////////////////////////////////
PLATFORM_IMPL(depthwise_conv2d, ENGINE_CUDA) {
auto input = INPUT_VARIABLE(0); // [bS, iH, iW, iC] (NHWC) or [bS, iC, iH, iW] (NCHW)
auto weights = INPUT_VARIABLE(1); // [kH, kW, iC, mC], [mC, iC, kH, kW], [mC, kH, kW, iC]
auto bias = block.width() > 2 ? INPUT_VARIABLE(2) : nullptr; // [oC] = iC*mC
auto output = OUTPUT_VARIABLE(0); // [bS, oH, oW, iC*mC] (NHWC) or [bS, iC*mC, oH, oW] (NCHW)
REQUIRE_TRUE(input->rankOf() == 4, 0,
"DEPTHWISECONV2D CUDNN OP: rank of input array must be equal to 4, but got %i instead !",
input->rankOf());
REQUIRE_TRUE(weights->rankOf() == 4, 0,
"DEPTHWISECONV2D CUDNN OP: rank of weights array must be equal to 4, but got %i instead !",
weights->rankOf());
LongType kH = INT_ARG(0) > 0 ? INT_ARG(0) : static_cast<LongType>(weights->sizeAt(0)); // filter(kernel) height
LongType kW = INT_ARG(1) > 0 ? INT_ARG(1) : static_cast<LongType>(weights->sizeAt(1)); // filter(kernel) width
LongType sH = INT_ARG(2); // strides height
LongType sW = INT_ARG(3); // strides width
LongType pH = INT_ARG(4); // paddings height
LongType pW = INT_ARG(5); // paddings width
LongType dH = INT_ARG(6); // dilations height
LongType dW = INT_ARG(7); // dilations width
int paddingMode = INT_ARG(8); // 0-VALID, 1-SAME
int isNCHW = block.getIArguments()->size() > 9 ? !INT_ARG(9) : 1; // INT_ARG(9): 0-NCHW, 1-NHWC
int wFormat = block.getIArguments()->size() > 10
? INT_ARG(10)
: 0; // 0 - [kH, kW, iC, mC], 1 - [mC, iC, kH, kW], 2 - [mC, kH, kW, iC]
LongType bS, iC, iH, iW, mC, oC, oH, oW; // batch size, input channels, input height/width, channels multiplier(oC =
// iC*mC), output channels, output height/width
LongType indIOioC, indIiH, indWmC, indWiC, indWkH, indOoH; // corresponding indexes
ConvolutionUtils::getSizesAndIndexesConv2d(isNCHW, wFormat, *input, *output, bS, iC, iH, iW, oC, oH, oW, indIOioC,
indIiH, indWiC, indWmC, indWkH, indOoH);
mC = weights->sizeAt(indWmC); // channels multiplier
ConvolutionUtils::calcPadding2D(pH, pW, oH, oW, iH, iW, kH, kW, sH, sW, dH, dW, paddingMode);
std::vector<LongType> expectedWeightsShape = ConvolutionUtils::expectWeightsShape(wFormat, kH, kW, iC, mC);
REQUIRE_TRUE(weights->isSameShape(expectedWeightsShape), 0,
"DEPTHWISECONV2D CUDNN OP: wrong shape of weights array, expected is %s, but got %s instead !",
ShapeUtils::shapeAsString(expectedWeightsShape).c_str(), ShapeUtils::shapeAsString(weights).c_str());
REQUIRE_TRUE(
output->sizeAt(indIOioC) == iC * mC, 0,
"DEPTHWISECONV2D CUDNN OP: the output_channels must be equal to input_channels * channels_multiplier = %i !",
iC * mC);
if (bias)
REQUIRE_TRUE(bias->rankOf() <= 2 && oC == bias->lengthOf(), 0,
"DEPTHWISECONV2D CUDNN OP: wrong shape of array with biases, expected rank, length: <=2, %i, but got "
"%i, %i instead !",
oC, bias->rankOf(), bias->lengthOf());
std::vector<LongType> wPermut; // cudnn support format {oC, iC/groupCount, kH, kW} only, mC = 1, oC = iC (groupCount ==
// iC) that is {iC, mC, kH, kW} in our case
if (0 == wFormat)
wPermut = {2, 3, 0, 1}; // kH, kW, iC, mC -> iC, mC, kH, kW
else if (1 == wFormat)
wPermut = {1, 0, 2, 3}; // mC, iC, kH, kW -> iC, mC, kH, kW
else
wPermut = {3, 0, 1, 2}; // mC, kH, kW, iC -> iC, mC, kH, kW
std::vector<sd::LongType > perm = {iC, mC, kH, kW};
NDArray * uNewWeights = new NDArray(weights->ordering(),perm, weights->dataType(), weights->getContext());
NDArray assign = weights->permute(wPermut,false,false);
uNewWeights->assign(&assign);
std::unique_ptr<NDArray> tmpInput = {};
if (paddingMode == 1) { // in same paddingMode cudnn doesn't support asymmetric left/right top/bottopm paddings
auto ret = checkConv2dCUDNNPadAsymmetric(input, nullptr, iH, iW, oH, oW, kH, kW, sH, sW, pH, pW, dH, dW, isNCHW);
tmpInput = std::move(std::get<0>(ret));
if (tmpInput) input = tmpInput.get();
}
depthwiseConv2dCUDNN(block.launchContext(), input, uNewWeights, bias, output, kH, kW, sH, sW, pH, pW, dH, dW,
paddingMode, isNCHW);
return Status::OK;
}
//////////////////////////////////////////////////////////////////////////
PLATFORM_CHECK(depthwise_conv2d, ENGINE_CUDA) {
auto input = INPUT_VARIABLE(0); // [bS, iH, iW, iC] (NHWC) or [bS, iC, iH, iW] (NCHW)
auto weights = INPUT_VARIABLE(1); // [kH, kW, iC, mC], [mC, iC, kH, kW], [mC, kH, kW, iC]
auto bias = block.width() > 2 ? INPUT_VARIABLE(2) : nullptr; // [oC] = iC*mC
const int paddingMode = INT_ARG(8); // 0-VALID, 1-SAME, 2-CAUSAL
const int wFormat = block.getIArguments()->size() > 10
? INT_ARG(10)
: 0; // 0 - [kH, kW, iC, mC], 1 - [mC, iC, kH, kW], 2 - [mC, kH, kW, iC]
Requirements req("CUDNN DEPTHWISE_CONV2d OP");
req.expectNotEq(makeInfoVariable(paddingMode, "paddingMode"), 2) &&
req.expectEq(makeInfoVariable(weights->sizeAt(0 == wFormat ? 3 : 0), "weights#mC"), 1) &&
req.expectIn(makeInfoVariable(input->dataType(), TYPE_MSG_INPUT0),
{HALF, FLOAT32, DOUBLE}) &&
req.expectIn(makeInfoVariable(weights->dataType(), TYPE_MSG_INPUT1),
{HALF, FLOAT32, DOUBLE});
if (bias) {
req.expectIn(makeInfoVariable(bias->dataType(), TYPE_MSG_INPUT_ "#bias"),
{HALF, FLOAT32, DOUBLE});
}
req.logTheSuccess();
return req;
}
//////////////////////////////////////////////////////////////////////////
PLATFORM_IMPL(depthwise_conv2d_bp, ENGINE_CUDA) {
auto input = INPUT_VARIABLE(0); // [bS, iH, iW, iC] (NDHWC) or [bS, iC, iH, iW] (NCDHW)
auto weights = INPUT_VARIABLE(1); // [kH, kW, iC, mC], [mC, iC, kH, kW], [mC, kH, kW, iC]
auto bias = block.width() > 3 ? INPUT_VARIABLE(2) : nullptr; // [oC] = [iC*mC]
auto gradO = block.width() > 3
? INPUT_VARIABLE(3)
: INPUT_VARIABLE(2); // [bS, oH, oW, oC] (NDHWC) or [bS, oC, oH, oW] (NCDHW), epsilon_next
auto gradI = OUTPUT_VARIABLE(0); // [bS, iH, iW, iC] (NDHWC) or [bS, iC, iH, iW] (NCDHW), epsilon
auto gradW = OUTPUT_VARIABLE(1); // [kH, kW, iC, mC], [mC, iC, kH, kW], [mC, kH, kW, iC]
auto gradB = block.width() > 3 ? OUTPUT_VARIABLE(2) : nullptr; // [oC]
REQUIRE_TRUE(input->rankOf() == 4, 0,
"DEPTHWISECONV2D_BP CUDNN OP: rank of input array must be equal to 4, but got %i instead !",
input->rankOf());
REQUIRE_TRUE(weights->rankOf() == 4, 0,
"DEPTHWISECONV2D_BP CUDNN OP: rank of weights array must be equal to 4, but got %i instead !",
weights->rankOf());
REQUIRE_TRUE(gradO->rankOf() == 4, 0,
"DEPTHWISECONV2D_BP CUDNN OP: rank of output gradients (next epsilon) array must be equal to 4, but got "
"%i instead !",
gradO->rankOf());
LongType kH = INT_ARG(0) > 0 ? INT_ARG(0) : static_cast<LongType>(weights->sizeAt(0)); // filter(kernel) height
LongType kW = INT_ARG(1) > 0 ? INT_ARG(1) : static_cast<LongType>(weights->sizeAt(1)); // filter(kernel) width
LongType sH = INT_ARG(2); // strides height
LongType sW = INT_ARG(3); // strides width
LongType pH = INT_ARG(4); // paddings height
LongType pW = INT_ARG(5); // paddings width
LongType dH = INT_ARG(6); // dilations height
LongType dW = INT_ARG(7); // dilations width
int paddingMode = INT_ARG(8); // 0-VALID, 1-SAME
int isNCHW = block.getIArguments()->size() > 9 ? !INT_ARG(9) : 1; // INT_ARG(9): 1-NHWC, 0-NCHW
int wFormat = block.getIArguments()->size() > 10
? INT_ARG(10)
: 0; // 0 - [kH, kW, iC, mC], 1 - [mC, iC, kH, kW], 2 - [mC, kH, kW, iC]
LongType bS, iC, iH, iW, mC, oC, oH, oW; // batch size, input channels, input height/width, channels multiplier(oC =
// iC*mC), output channels, output height/width
LongType indIOioC, indIiH, indWmC, indWiC, indWkH, indOoH; // corresponding indexes
ConvolutionUtils::getSizesAndIndexesConv2d(isNCHW, wFormat, *input, *gradO, bS, iC, iH, iW, oC, oH, oW, indIOioC,
indIiH, indWiC, indWmC, indWkH, indOoH);
mC = weights->sizeAt(indWmC); // channels multiplier
LongType trueoH, trueoW; // correct output height, width
ConvolutionUtils::calcOutSizePool2D(trueoH, trueoW, kH, kW, sH, sW, pH, pW, dH, dW, iH, iW, paddingMode);
ConvolutionUtils::calcPadding2D(pH, pW, oH, oW, iH, iW, kH, kW, sH, sW, dH, dW, paddingMode);
std::vector<LongType> expectedGradOShape =
ShapeUtils::composeShapeUsingDimsAndIdx({bS, oC, trueoH, trueoW, 0, indIOioC, indOoH, indOoH + 1});
std::vector<LongType> expectedWeightsShape = ConvolutionUtils::expectWeightsShape(wFormat, kH, kW, iC, mC);
REQUIRE_TRUE(gradO->isSameShape(expectedGradOShape), 0,
"DEPTHWISECONV2D_BP CUDNN OP: wrong shape of output gradients (next epsilon) array, expected is %s, but "
"got %s instead !",
ShapeUtils::shapeAsString(expectedGradOShape).c_str(), ShapeUtils::shapeAsString(gradO).c_str());
REQUIRE_TRUE(weights->isSameShape(expectedWeightsShape), 0,
"DEPTHWISECONV2D_BP CUDNN OP: wrong shape of weights array, expected is %s, but got %s instead !",
ShapeUtils::shapeAsString(expectedWeightsShape).c_str(), ShapeUtils::shapeAsString(weights).c_str());
if (bias)
REQUIRE_TRUE(bias->rankOf() <= 2 && oC == bias->lengthOf(), 0,
"DEPTHWISECONV2D_BP CUDNN OP: wrong shape of array with biases, expected rank, length: <=2, %i, but "
"got %i, %i instead !",
oC, bias->rankOf(), bias->lengthOf());
std::vector<LongType> wPermut, gradWPermut; // cudnn support format {oC, iC/groupCount, kH, kW} only, mC = 1, oC = iC
// (groupCount == iC) that is {iC, mC, kH, kW}
if (0 == wFormat) {
wPermut = {2, 3, 0, 1}; // kH, kW, iC, mC -> iC, mC, kH, kW
gradWPermut = {2, 3, 0, 1}; // iC, mC, kH, kW -> kH, kW, iC, mC
} else if (1 == wFormat) {
wPermut = {1, 0, 2, 3}; // mC, iC, kH, kW -> iC, mC, kH, kW
gradWPermut = {1, 0, 2, 3}; // iC, mC, kH, kW -> mC, iC, kH, kW
} else {
wPermut = {3, 0, 1, 2}; // mC, kH, kW, iC -> iC, mC, kH, kW
gradWPermut = {1, 2, 3, 0}; // iC, mC, kH, kW -> mC, kH, kW, iC
}
std::unique_ptr<NDArray> tmpGradI = {}, tmpInput = {};
std::vector<sd::LongType> shape = {iC, mC, kH, kW};
NDArray * uNewGradW =
new NDArray(gradW->ordering(),shape, gradW->dataType(), gradW->getContext());
NDArray * uNewWeights =
new NDArray(weights->ordering(),shape, weights->dataType(), weights->getContext());
NDArray assign = weights->permute(wPermut,false,false);
uNewWeights->assign(&assign);
NDArray* newInput = input;
NDArray* newGradI = gradI;
if (paddingMode == 1) { // in same paddingMode cudnn doesn't support asymmetric left/right top/bottopm paddings
auto ret = checkConv2dCUDNNPadAsymmetric(input, gradI, iH, iW, oH, oW, kH, kW, sH, sW, pH, pW, dH, dW, isNCHW);
tmpInput = std::move(std::get<0>(ret));
tmpGradI = std::move(std::get<1>(ret));
if (tmpInput) newInput = tmpInput.get();
if (tmpGradI) newGradI = tmpGradI.get();
}
depthwiseConv2dBpCUDNN(block.launchContext(), newInput, uNewWeights, gradO, newGradI, uNewGradW, gradB,
kH, kW, sH, sW, pH, pW, dH, dW, paddingMode, isNCHW);
uNewGradW->permutei(gradWPermut,false,false);
gradW->assign(uNewGradW);
if (newInput != input) {
if (isNCHW) {
NDArray assign = (*newGradI)({0, 0, 0, 0, 0, gradI->sizeAt(2), 0, gradI->sizeAt(3)});
gradI->assign(&assign);
} else {
NDArray assign = (*newGradI)({0, 0, 0, gradI->sizeAt(1), 0, gradI->sizeAt(2), 0, 0});
gradI->assign(&assign);
}
}
return Status::OK;
}
PLATFORM_CHECK(depthwise_conv2d_bp, ENGINE_CUDA) {
auto input = INPUT_VARIABLE(0); // [bS, iH, iW, iC] (NDHWC) or [bS, iC, iH, iW] (NCDHW)
auto weights = INPUT_VARIABLE(1); // [kH, kW, iC, mC], [mC, iC, kH, kW], [mC, kH, kW, iC]
auto bias = block.width() > 3 ? INPUT_VARIABLE(2) : nullptr; // [oC] = [iC*mC]
auto gradO = block.width() > 3
? INPUT_VARIABLE(3)
: INPUT_VARIABLE(2); // [bS, oH, oW, oC] (NDHWC) or [bS, oC, oH, oW] (NCDHW), epsilon_next
const int paddingMode = INT_ARG(8); // 0-VALID, 1-SAME, 2-CAUSAL
const int isNCHW = block.getIArguments()->size() > 9 ? !INT_ARG(9) : 1; // INT_ARG(9): 0-NCHW, 1-NHWC
const int wFormat = block.getIArguments()->size() > 10
? INT_ARG(10)
: 0; // 0 - [kH, kW, iC, mC], 1 - [mC, iC, kH, kW], 2 - [mC, kH, kW, iC]
Requirements req("CUDNN DEPTHWISE_CONV2d_BP OP");
const auto inType = input->dataType();
const auto wType = weights->dataType();
const auto gType = gradO->dataType();
req.expectNotEq(makeInfoVariable(paddingMode, "paddingMode"), 2) &&
req.expectTrue(makeInfoVariable(isNCHW, "isNCHW")) &&
req.expectEq(makeInfoVariable(weights->sizeAt(0 == wFormat ? 3 : 0), "weights#mC"), 1) &&
req.expectIn(makeInfoVariable(input->dataType(), TYPE_MSG_INPUT0),
{HALF, FLOAT32, DOUBLE}) &&
req.expectIn(makeInfoVariable(weights->dataType(), TYPE_MSG_INPUT1),
{HALF, FLOAT32, DOUBLE});
if (bias) {
req.expectIn(makeInfoVariable(bias->dataType(), TYPE_MSG_INPUT_ "#bias"),
{HALF, FLOAT32, DOUBLE}) &&
req.expectIn(makeInfoVariable(gradO->dataType(), TYPE_MSG_INPUT3),
{HALF, FLOAT32, DOUBLE});
} else {
req.expectIn(makeInfoVariable(gradO->dataType(), TYPE_MSG_INPUT2),
{HALF, FLOAT32, DOUBLE});
}
req.logTheSuccess();
return req;
}
} // namespace platforms
} // namespace ops
} // namespace sd
@@ -0,0 +1,673 @@
/* ******************************************************************************
*
*
* 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 AbdelRauf
//
#include <array/NDArrayFactory.h>
#include <ops/declarable/OpRegistrator.h>
#include "cudnnUtils.h"
namespace sd {
namespace ops {
namespace platforms {
// our implementation designed for 1 physical layer
constexpr int numLayers = 1;
// we will copy without using cudnnGetRNNLinLayerMatrixParams : 1 pseudo layer , isBidirectional : 2 pseudo layer
void copyWeights(const cudaStream_t &stream, bool isBidirectional, uint8_t *weightsSpace, size_t weightsSize,
uint8_t *inputWeightsData, uint8_t *recurrentWeightsData, uint8_t *biasesData, LongType inputSize,
int hiddenSize, int dataTypeSize) {
int pseudo_layer_count = isBidirectional ? 2 : 1;
uint8_t *wptr = weightsSpace;
auto wEnd = wptr + weightsSize;
// copy size for 1 full pseudo layer
// in bidirectional 1 layer consist of 2 pseduo layers
auto input_pseudo_size = 4 * inputSize * hiddenSize * dataTypeSize;
auto hidden_pseudo_size = 4 * hiddenSize * hiddenSize * dataTypeSize;
for (LongType i = 0; i < pseudo_layer_count; i++) {
if (wptr + input_pseudo_size + hidden_pseudo_size > wEnd) return;
// copy input weights
if (inputWeightsData) {
cudaMemcpyAsync(wptr, inputWeightsData, input_pseudo_size, cudaMemcpyDeviceToDevice, stream);
inputWeightsData += input_pseudo_size;
}
wptr += input_pseudo_size;
// copy recurrent weights
if (recurrentWeightsData) {
cudaMemcpyAsync(wptr, recurrentWeightsData, hidden_pseudo_size, cudaMemcpyDeviceToDevice, stream);
recurrentWeightsData += hidden_pseudo_size;
}
wptr += hidden_pseudo_size;
}
// copy bias first 4
auto bias_size = 4 * hiddenSize * dataTypeSize;
for (int i = 0; i < pseudo_layer_count; i++) {
// refill first 4 biases
if (biasesData && wptr + bias_size < wEnd) {
cudaMemcpyAsync(wptr, biasesData, bias_size, cudaMemcpyDeviceToDevice, stream);
biasesData += bias_size;
}
wptr += bias_size;
// refill next 4 with zeros
if (wptr + bias_size < wEnd) {
cudaMemsetAsync(wptr, 0, bias_size, stream);
wptr += bias_size;
}
}
// memset the rest
if (wEnd - wptr) cudaMemsetAsync(wptr, 0, wEnd - wptr, stream);
}
void cudnn_rnn_old(LaunchContext *contextPtr, int dataFormat, NDArray *input, NDArray *inputWeights,
NDArray *recurrentWeights, NDArray *biases, NDArray *prevAct, NDArray *prevMemCell,
NDArray *outputActivations, NDArray *finalTimeStepActivations, NDArray *finalMemCellState,
LongType maxSeqLength, LongType batchSize, LongType inputSize, LongType hiddenSize, double cellClip,
bool isBidirectional) {
sd_debug("cudnn rnn api %s \n", "v6");
bool training = false;
cudnnHandle_t handle = *(reinterpret_cast<cudnnHandle_t *>(contextPtr->getCuDnnHandle()));
auto stream = *(contextPtr->getCudaStream());
CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnSetStream), cudnnSetStream(handle, stream));
CudnnTensorList xDescList(maxSeqLength);
CudnnTensorList yDescList(maxSeqLength);
auto cudnnType = cudnnDataType(input->dataType());
auto dataTypeSize = input->sizeOfT();
CudnnTensor hxDesc, cxDesc, hyDesc, cyDesc;
constexpr int rankOf = 3;
const int numDirections = isBidirectional ? 2 : 1;
const int dimsX[rankOf] = {static_cast<int>(batchSize), static_cast<int>(inputSize), 1};
const int stridesX[rankOf] = {static_cast<int>(inputSize), 1, 1};
const int dimsY[rankOf] = {static_cast<int>(batchSize), static_cast<int>(hiddenSize * numDirections), 1};
const int stridesY[rankOf] = {static_cast<int>(hiddenSize * numDirections), 1, 1};
const int dimC[rankOf] = {static_cast<int>(numLayers * numDirections), static_cast<int>(batchSize), static_cast<int>(hiddenSize)};
const int strideC[rankOf] = {static_cast<int>(batchSize * hiddenSize), static_cast<int>(hiddenSize), 1};
for (int i = 0; i < maxSeqLength; i++) {
xDescList.set(i, cudnnType, rankOf, dimsX, stridesX);
yDescList.set(i, cudnnType, rankOf, dimsY, stridesY);
}
auto xDesc0 = xDescList.get(0);
hxDesc.set(cudnnType, rankOf, dimC, strideC);
cxDesc.set(cudnnType, rankOf, dimC, strideC);
hyDesc.set(cudnnType, rankOf, dimC, strideC);
cyDesc.set(cudnnType, rankOf, dimC, strideC);
PointersManager manager(contextPtr, __func__);
// dropout section
DropoutDesc dropoutDesc(nullptr);
// dropout
float dropout = 0;
size_t sizeInBytes = 0;
void *droupoutMem = nullptr;
uint64_t seed = 1; // seed
if (dropout != 0) {
dropoutDesc.create();
CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnDropoutGetStatesSize), cudnnDropoutGetStatesSize(handle, &sizeInBytes));
// allocate and set
droupoutMem = manager.allocateDevMem(sizeInBytes);
dropoutDesc.set(handle, dropout, droupoutMem, sizeInBytes, seed);
}
// RNN
RnnDesc rnnDesc;
cudnnRNNMode_t rnnCellMode = CUDNN_LSTM;
cudnnRNNAlgo_t algo = CUDNN_RNN_ALGO_STANDARD;
auto direction = isBidirectional ? CUDNN_BIDIRECTIONAL : CUDNN_UNIDIRECTIONAL;
auto mathPrec = cudnnType;
// Note: We will set some parameters manually
constexpr auto inputMode = CUDNN_LINEAR_INPUT;
rnnDesc.setUsingOldAPI(handle, inputMode, direction, rnnCellMode, algo, mathPrec, hiddenSize, numLayers, dropoutDesc);
#if CUDNN_VERSION >= CUDNN_CLIPPING_API_VER
if (cellClip > 0 && cudnnGetVersion() >= CUDNN_CLIPPING_API_VER) {
CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnRNNSetClip), cudnnRNNSetClip(handle, rnnDesc, CUDNN_RNN_CLIP_MINMAX,
CUDNN_PROPAGATE_NAN, -cellClip, cellClip));
}
#endif
// set up parameters
size_t weightsSize = 0;
CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnGetRNNParamsSize),
cudnnGetRNNParamsSize(handle, rnnDesc, xDesc0, &weightsSize, cudnnType));
FilterDesc wDesc;
int dimW[] = {static_cast<int>(weightsSize / dataTypeSize), 1, 1};
wDesc.set(cudnnType, CUDNN_TENSOR_NCHW, 3, dimW);
// allocation
void *weightsSpace = manager.allocateDevMem(weightsSize);
size_t workSpaceSizeInBytes = 0;
size_t reserveSpaceSizeInBytes = 0;
CHECK_CUDNN_FAILURE_MSG(
STRINGIZE(cudnnGetRNNWorkspaceSize),
cudnnGetRNNWorkspaceSize(handle, rnnDesc, maxSeqLength, xDescList.getDescriptors(), &workSpaceSizeInBytes));
void *workSpace = manager.allocateDevMem(workSpaceSizeInBytes);
void *reserveSpace = nullptr;
// training
if (training) {
CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnGetRNNTrainingReserveSize),
cudnnGetRNNTrainingReserveSize(handle, rnnDesc, maxSeqLength, xDescList.getDescriptors(),
&reserveSpaceSizeInBytes));
reserveSpace = manager.allocateDevMem(reserveSpaceSizeInBytes);
}
NDArray::prepareSpecialUse({outputActivations, finalTimeStepActivations, finalMemCellState},
{input, inputWeights, recurrentWeights, biases, prevAct, prevMemCell});
uint8_t *biasesData = biases ? (uint8_t *)biases->specialBuffer() : nullptr;
auto prevActData = prevAct ? prevAct->specialBuffer() : nullptr;
auto prevMemCellData = prevMemCell ? prevMemCell->specialBuffer() : nullptr;
auto finalTimeStepActivationsData = finalTimeStepActivations ? finalTimeStepActivations->specialBuffer() : nullptr;
auto finalMemCellStateData = finalMemCellState ? finalMemCellState->specialBuffer() : nullptr;
// dimension 4*nOut implies order it, ft, c't, ot
// input gate, forget gate, new gate, output gate, input gate, forget gate, new gate, output gate
// Note: our weights should be transposed and duplicated with C order to match cudnn ones
NDArray inputWeightsT, recurrentWeightsT;
uint8_t *inputWeightsData = nullptr;
uint8_t *recurrentWeightsData = nullptr;
if (inputWeights) {
inputWeightsT =
inputWeights->rankOf() == 3 ? inputWeights->permute({0, 2, 1}, 0, false).dup('c') : inputWeights->transpose().dup('c');
inputWeightsData = (uint8_t *)inputWeightsT.specialBuffer();
}
if (recurrentWeights) {
recurrentWeightsT = recurrentWeights->rankOf() == 3 ? recurrentWeights->permute({0, 2, 1}, 0, false).dup('c')
: recurrentWeights->transpose().dup('c');
recurrentWeightsData = (uint8_t *)recurrentWeightsT.specialBuffer();
}
// copy without cudnnGetRNNLinLayerMatrixParams
copyWeights(stream, isBidirectional, (uint8_t *)weightsSpace, weightsSize, inputWeightsData, recurrentWeightsData,
biasesData, inputSize, hiddenSize, dataTypeSize);
// permute based on dataformat
NDArray *argX = input;
NDArray *argOutput = outputActivations;
NDArray permutedX, outputH;
if (outputActivations != nullptr && (dataFormat != 0 || outputActivations->ordering() != 'c')) {
outputH = NDArray('c', std::vector<LongType>{maxSeqLength, batchSize, (numDirections * hiddenSize)},
outputActivations->dataType(), contextPtr);
argOutput = &outputH;
}
if (dataFormat == 1) {
permutedX = input->permute({1, 0, 2}, 0, false).dup('c');
argX = &permutedX;
}
auto xData = argX->specialBuffer();
auto yData = argOutput ? argOutput->specialBuffer() : nullptr;
if (training) {
CHECK_CUDNN_FAILURE_MSG(
STRINGIZE(cudnnRNNForwardTraining),
cudnnRNNForwardTraining(handle, rnnDesc, (int)maxSeqLength, xDescList.getDescriptors(), xData, hxDesc,
prevActData, cxDesc, prevMemCellData, wDesc, weightsSpace, yDescList.getDescriptors(),
yData, hyDesc, finalTimeStepActivationsData, cyDesc, finalMemCellStateData, workSpace,
workSpaceSizeInBytes, reserveSpace, reserveSpaceSizeInBytes));
} else {
CHECK_CUDNN_FAILURE_MSG(
STRINGIZE(cudnnRNNForwardInference),
cudnnRNNForwardInference(handle, rnnDesc, (int)maxSeqLength, xDescList.getDescriptors(), xData, hxDesc,
prevActData, cxDesc, prevMemCellData, wDesc, weightsSpace, yDescList.getDescriptors(),
yData, hyDesc, finalTimeStepActivationsData, cyDesc, finalMemCellStateData, workSpace,
workSpaceSizeInBytes));
}
// remap output
if (outputActivations != nullptr && argOutput != outputActivations) {
// refill output
if (dataFormat == 1) {
std::vector<sd::LongType> permute = {1,0,2};
NDArray assign = argOutput->permute(permute, 0, false);
outputActivations->assign(&assign);
}
}
NDArray::registerSpecialUse({outputActivations, finalTimeStepActivations, finalMemCellState},
{input, inputWeights, recurrentWeights, biases, prevAct, prevMemCell});
return;
}
#if CUDNN_VERSION >= CUDNN_NEW_RNN_API_VER
void cudnn_rnn_v8(LaunchContext *contextPtr, int dataFormat, NDArray *input, NDArray *seqLengthArray,
NDArray *inputWeights, NDArray *recurrentWeights, NDArray *biases, NDArray *prevAct,
NDArray *prevMemCell, NDArray *outputActivations, NDArray *finalTimeStepActivations,
NDArray *finalMemCellState, int maxSeqLength, int batchSize, int inputSize, int hiddenSize,
double cellClip, bool isBidirectional) {
sd_debug("cudnn rnn api %s \n", "v8");
// seqLengthArray should be int
NDArray *argSeqNdArray = nullptr;
NDArray seqArrIntData;
if (seqLengthArray) {
if (seqLengthArray->ews() == 1 && seqLengthArray->dataType() == INT32) {
argSeqNdArray = seqLengthArray;
} else {
if (seqLengthArray->dataType() != INT32) {
seqArrIntData = seqLengthArray->cast(INT32);
if (seqArrIntData.ews() != 1) seqArrIntData = seqArrIntData.dup('c');
} else {
seqArrIntData = seqLengthArray->dup('c');
}
argSeqNdArray = &seqArrIntData;
}
} else {
seqArrIntData = NDArray('c', std::vector<LongType>{batchSize}, INT32, contextPtr);
seqArrIntData.assign(maxSeqLength);
argSeqNdArray = &seqArrIntData;
}
PointersManager manager(contextPtr, __func__);
bool training = false;
cudnnHandle_t handle = *(reinterpret_cast<cudnnHandle_t *>(contextPtr->getCuDnnHandle()));
auto stream = *(contextPtr->getCudaStream());
CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnSetStream), cudnnSetStream(handle, stream));
auto cudnnType = cudnnDataType(input->dataType());
auto dataTypeSize = input->sizeOfT();
CudnnTensor hDesc, cDesc;
constexpr int rankOf = 3;
const int numDirections = isBidirectional ? 2 : 1;
const int dimC[rankOf] = {numLayers * numDirections, batchSize, hiddenSize};
const int strideC[rankOf] = {batchSize * hiddenSize, hiddenSize, 1};
hDesc.set(cudnnType, rankOf, dimC, strideC);
cDesc.set(cudnnType, rankOf, dimC, strideC);
// dropout section
DropoutDesc dropoutDesc(nullptr);
// dropout
float dropout = 0;
size_t sizeInBytes = 0;
void *droupoutMem = nullptr;
uint64_t seed = 1; // seed
if (dropout != 0) {
dropoutDesc.create();
CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnDropoutGetStatesSize), cudnnDropoutGetStatesSize(handle, &sizeInBytes));
// allocate and set
droupoutMem = manager.allocateDevMem(sizeInBytes);
dropoutDesc.set(handle, dropout, droupoutMem, sizeInBytes, seed);
}
// RNN
RnnDesc rnnDesc;
cudnnRNNMode_t rnnCellMode = CUDNN_LSTM;
cudnnRNNAlgo_t algo = CUDNN_RNN_ALGO_STANDARD;
auto direction = isBidirectional ? CUDNN_BIDIRECTIONAL : CUDNN_UNIDIRECTIONAL;
auto mathPrec = cudnnType;
// Note: We will set some parameters manually. Some of them could be parameter in future
constexpr auto inputMode = CUDNN_LINEAR_INPUT;
bool use_tensor_ops = false; // could be parameter in future
#if CUDNN_VERSION >= CUDNN_NEW_RNN_API_VER
cudnnMathType_t mathType = use_tensor_ops ? CUDNN_TENSOR_OP_MATH : CUDNN_FMA_MATH;
#else
cudnnMathType_t mathType = use_tensor_ops ? CUDNN_TENSOR_OP_MATH : CUDNN_DEFAULT_MATH;
#endif
// disable projection
int projSize = hiddenSize;
cudnnRNNBiasMode_t bias_mode = CUDNN_RNN_DOUBLE_BIAS;
uint32_t aux_flags = CUDNN_RNN_PADDED_IO_ENABLED;
rnnDesc.set(algo, rnnCellMode, bias_mode, direction, inputMode, cudnnType, mathPrec, mathType, inputSize, hiddenSize,
projSize, numLayers, dropoutDesc, aux_flags);
if (cellClip > 0) {
CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnRNNSetClip), cudnnRNNSetClip(handle, rnnDesc, CUDNN_RNN_CLIP_MINMAX,
CUDNN_PROPAGATE_NAN, -cellClip, cellClip));
}
// set Data desc
RnnDataDesc xDataDesc, yDataDesc;
bool time_major = false;
float padding_fill = 0.0f;
auto hostSeqArr = bufferInHost<int>(*argSeqNdArray);
cudnnRNNDataLayout_t layout =
dataFormat == 0 ? CUDNN_RNN_DATA_LAYOUT_SEQ_MAJOR_UNPACKED : CUDNN_RNN_DATA_LAYOUT_BATCH_MAJOR_UNPACKED;
xDataDesc.set(cudnnType, layout, maxSeqLength, batchSize, inputSize, hostSeqArr, (void *)&padding_fill);
yDataDesc.set(cudnnType, layout, maxSeqLength, batchSize, hiddenSize * numDirections, hostSeqArr,
(void *)&padding_fill);
// set up parameters
size_t weightsSize = 0;
CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnGetRNNWeightSpaceSize),
cudnnGetRNNWeightSpaceSize(handle, rnnDesc, &weightsSize));
// allocation
void *weightsSpace = manager.allocateDevMem(weightsSize);
// Set up work space and reserved memory
void *workSpace = nullptr;
void *reserveSpace = nullptr;
size_t workSpaceSizeInBytes = 0;
size_t reserveSpaceSizeInBytes = 0;
cudnnForwardMode_t fwdMode = training ? CUDNN_FWD_MODE_TRAINING : CUDNN_FWD_MODE_INFERENCE;
CHECK_CUDNN_FAILURE_MSG(
STRINGIZE(cudnnGetRNNTempSpaceSizes),
cudnnGetRNNTempSpaceSizes(handle, rnnDesc, fwdMode, xDataDesc, &workSpaceSizeInBytes, &reserveSpaceSizeInBytes));
workSpace = manager.allocateDevMem(workSpaceSizeInBytes);
// training
if (training) {
reserveSpace = manager.allocateDevMem(reserveSpaceSizeInBytes);
}
NDArray::prepareSpecialUse({outputActivations, finalTimeStepActivations, finalMemCellState},
{input, inputWeights, recurrentWeights, biases, prevAct, prevMemCell, argSeqNdArray});
auto xData = input->specialBuffer();
uint8_t *biasesData = biases ? (uint8_t *)biases->specialBuffer() : nullptr;
auto prevActData = prevAct ? prevAct->specialBuffer() : nullptr;
auto prevMemCellData = prevMemCell ? prevMemCell->specialBuffer() : nullptr;
auto yData = outputActivations ? outputActivations->specialBuffer() : nullptr;
auto finalTimeStepActivationsData = finalTimeStepActivations ? finalTimeStepActivations->specialBuffer() : nullptr;
auto finalMemCellStateData = finalMemCellState ? finalMemCellState->specialBuffer() : nullptr;
// dimension 4*nOut implies order it, ft, c't, ot
// input gate, forget gate, new gate, output gate, input gate, forget gate, new gate, output gate
// Note: our weights should be transposed and duplicated with C order to match cudnn ones
NDArray inputWeightsT, recurrentWeightsT;
uint8_t *inputWeightsData = nullptr;
uint8_t *recurrentWeightsData = nullptr;
if (inputWeights) {
inputWeightsT =
inputWeights->rankOf() == 3 ? inputWeights->permute({0, 2, 1}).dup('c') : inputWeights->transpose().dup('c');
inputWeightsData = (uint8_t *)inputWeightsT.specialBuffer();
}
if (recurrentWeights) {
recurrentWeightsT = recurrentWeights->rankOf() == 3 ? recurrentWeights->permute({0, 2, 1}).dup('c')
: recurrentWeights->transpose().dup('c');
recurrentWeightsData = (uint8_t *)recurrentWeightsT.specialBuffer();
}
// copy without cudnnGetRNNLinLayerMatrixParams
copyWeights(stream, isBidirectional, (uint8_t *)weightsSpace, weightsSize, inputWeightsData, recurrentWeightsData,
biasesData, inputSize, hiddenSize, dataTypeSize);
CHECK_CUDNN_FAILURE_MSG(
STRINGIZE(cudnnRNNForward),
cudnnRNNForward(handle, rnnDesc, fwdMode, (const int32_t *)argSeqNdArray->specialBuffer(), xDataDesc, xData,
yDataDesc, yData, hDesc, prevActData, finalTimeStepActivationsData, cDesc, prevMemCellData,
finalMemCellStateData, weightsSize, weightsSpace, workSpaceSizeInBytes, workSpace,
reserveSpaceSizeInBytes, reserveSpace));
NDArray::registerSpecialUse({outputActivations, finalTimeStepActivations, finalMemCellState},
{input, inputWeights, recurrentWeights, biases, prevAct, prevMemCell});
return;
}
#endif
//////////////////////////////////////////////////////////////////////////
PLATFORM_IMPL(lstmLayer, ENGINE_CUDA) {
const auto dataFormat = INT_ARG(0); // for unidirectional: 0 = [sL, bS, nIn], 1 = [bS, sL ,nIn], 2 = [bS, nIn, sL],
// for bidirectional: 3 = [sL, 2, bS, nOut] (for ONNX)
const LongType directionMode =
INT_ARG(1); // direction: 0 = fwd, 1 = bwd, 2 = bidirectional sum, 3 = bidirectional concat, 4 = bidirectional
// extra output dim (in conjunction with format dataFormat = 3)
const auto hasBiases = B_ARG(0); // indicates whether biases array is provided
const auto hasSeqLenArray = B_ARG(1); // indicates whether seqLen array is provided
const auto hasInitH = B_ARG(2); // indicates whether initial output is provided
const auto hasInitC = B_ARG(3); // indicates whether initial cell state is provided
const auto hasPH = B_ARG(4); // indicates whether peephole connections are present
const auto retFullSeq = B_ARG(5); // indicates whether to return whole time sequence h {h_0, h_1, ... , h_sL-1}
const auto retLastH = B_ARG(6); // indicates whether to return output at last time step only, in this case shape
// would be [bS, nOut] (exact shape depends on dataFormat argument)
const auto retLastC = B_ARG(7); // indicates whether to return cells state at last time step only, in this case shape
// would be [bS, nOut] (exact shape depends on dataFormat argument)
const auto cellClip = T_ARG(0); // cell clipping value, if it = 0 then do not apply clipping
const auto x = INPUT_VARIABLE(0); // input
const auto Wx = INPUT_VARIABLE(1); // input weights
const auto Wr = INPUT_VARIABLE(2); // recurrent weights
int count = 3;
const auto b = hasBiases ? INPUT_VARIABLE(count++) : nullptr; // biases
const auto seqLengthArray = hasSeqLenArray ? INPUT_VARIABLE(count++) : nullptr; // seqLen vector
const auto hI = hasInitH ? INPUT_VARIABLE(count++) : nullptr; // initial output
const auto cI = hasInitC ? INPUT_VARIABLE(count++) : nullptr; // initial cell state
const auto Wp = hasPH ? INPUT_VARIABLE(count++) : nullptr; // peephole weights
count = 0;
auto h = retFullSeq ? OUTPUT_VARIABLE(count++) : nullptr; // output
auto hL = retLastH ? OUTPUT_VARIABLE(count++) : nullptr; // output at last step
auto cL = retLastC ? OUTPUT_VARIABLE(count++) : nullptr; // cell state at last step
REQUIRE_TRUE(cellClip >= 0, 0, "LSTM_LAYER operation: cell clipping value should be nonnegative (>=0) !");
REQUIRE_TRUE(retFullSeq || retLastH || retLastC, 0,
"LSTM_LAYER operation: please specify what output arrays to produce !");
// evaluate dimensions
const LongType seqLength = dataFormat == 3 ? x->sizeAt(0) : x->sizeAt(dataFormat);
const LongType bS = dataFormat == 1 || dataFormat == 2 ? x->sizeAt(0) : x->sizeAt(1);
const LongType nIn = dataFormat == 2 ? x->sizeAt(1) : x->sizeAt(2);
const LongType nOut = Wx->sizeAt(-1) / 4;
const LongType hiddenSize = nOut;
auto contextPtr = block.launchContext();
bool isBidirectional = directionMode >= 2;
if (!isBidirectional) { // no bidirectional
// Wx validation
if (Wx->rankOf() != 2 || Wx->sizeAt(0) != nIn)
REQUIRE_TRUE(false, 0, "LSTM_LAYER operation: wrong shape of input weights, expected is %s, but got %s instead !",
ShapeUtils::shapeAsString({nIn, 4 * nOut}).c_str(), ShapeUtils::shapeAsString(Wx).c_str());
// Wr validation
if (Wr->rankOf() != 2 || Wr->sizeAt(0) != nOut || Wr->sizeAt(1) != 4 * nOut)
REQUIRE_TRUE(false, 0,
"LSTM_LAYER operation: wrong shape of recurrent weights, expected is %s, but got %s instead !",
ShapeUtils::shapeAsString({nOut, 4 * nOut}).c_str(), ShapeUtils::shapeAsString(Wr).c_str());
// biases validation
if (b != nullptr && (b->rankOf() != 1 || b->sizeAt(0) != 4 * nOut))
REQUIRE_TRUE(false, 0, "LSTM_LAYER operation: wrong shape of biases, expected is %s, but got %s instead !",
ShapeUtils::shapeAsString({4 * nOut}).c_str(), ShapeUtils::shapeAsString(b).c_str());
// initial output validation
if (hI != nullptr && (hI->rankOf() != 2 || hI->sizeAt(0) != bS || hI->sizeAt(1) != nOut))
REQUIRE_TRUE(false, 0,
"LSTM_LAYER operation: wrong shape of initial output, expected is %s, but got %s instead !",
ShapeUtils::shapeAsString({bS, nOut}).c_str(), ShapeUtils::shapeAsString(hI).c_str());
// initial cell validation
if (cI != nullptr && (cI->rankOf() != 2 || cI->sizeAt(0) != bS || cI->sizeAt(1) != nOut))
REQUIRE_TRUE(false, 0,
"LSTM_LAYER operation: wrong shape of initial cell state, expected is %s, but got %s instead !",
ShapeUtils::shapeAsString({bS, nOut}).c_str(), ShapeUtils::shapeAsString(cI).c_str());
} else { // bidirectional
// Wx validation
if (Wx->rankOf() != 3 || Wx->sizeAt(0) != 2 || Wx->sizeAt(1) != nIn)
REQUIRE_TRUE(false, 0, "LSTM_LAYER operation: wrong shape of input weights, expected is %s, but got %s instead !",
ShapeUtils::shapeAsString({2, nIn, 4 * nOut}).c_str(), ShapeUtils::shapeAsString(Wx).c_str());
// Wr validation
if (Wr->rankOf() != 3 || Wr->sizeAt(0) != 2 || Wr->sizeAt(1) != nOut || Wr->sizeAt(2) != 4 * nOut)
REQUIRE_TRUE(false, 0,
"LSTM_LAYER operation: wrong shape of recurrent weights, expected is %s, but got %s instead !",
ShapeUtils::shapeAsString({2, nOut, 4 * nOut}).c_str(), ShapeUtils::shapeAsString(Wr).c_str());
// biases validation
if (b != nullptr && (b->rankOf() != 2 || b->sizeAt(0) != 2 || b->sizeAt(1) != 4 * nOut))
REQUIRE_TRUE(false, 0, "LSTM_LAYER operation: wrong shape of biases, expected is %s, but got %s instead !",
ShapeUtils::shapeAsString({2, 4 * nOut}).c_str(), ShapeUtils::shapeAsString(b).c_str());
// initial output validation
if (hI != nullptr && (hI->rankOf() != 3 || hI->sizeAt(0) != 2 || hI->sizeAt(1) != bS || hI->sizeAt(2) != nOut))
REQUIRE_TRUE(false, 0,
"LSTM_LAYER operation: wrong shape of initial output, expected is %s, but got %s instead !",
ShapeUtils::shapeAsString({2, bS, nOut}).c_str(), ShapeUtils::shapeAsString(hI).c_str());
// initial cell validation
if (cI != nullptr && (cI->rankOf() != 3 || cI->sizeAt(0) != 2 || cI->sizeAt(1) != bS || cI->sizeAt(2) != nOut))
REQUIRE_TRUE(false, 0,
"LSTM_LAYER operation: wrong shape of initial cell state, expected is %s, but got %s instead !",
ShapeUtils::shapeAsString({2, bS, nOut}).c_str(), ShapeUtils::shapeAsString(cI).c_str());
}
#if CUDNN_VERSION < CUDNN_NEW_RNN_API_VER
cudnn_rnn_old(contextPtr, dataFormat, x, Wx, Wr, b, hI, cI, h, hL, cL, seqLength, bS, nIn, hiddenSize,
(double)cellClip, isBidirectional);
#else
if (cudnnGetVersion() >= CUDNN_NEW_RNN_API_VER) {
cudnn_rnn_v8(contextPtr, dataFormat, x, seqLengthArray, Wx, Wr, b, hI, cI, h, hL, cL, seqLength, bS, nIn,
hiddenSize, (double)cellClip, isBidirectional);
} else {
cudnn_rnn_old(contextPtr, dataFormat, x, Wx, Wr, b, hI, cI, h, hL, cL, seqLength, bS, nIn, hiddenSize,
(double)cellClip, isBidirectional);
}
#endif
return Status::OK;
}
// Cudnn Lstm:
// Forward inference implemented using v6, and v8 (when version > 8.0.1) api calls.
// As our Cuda Lstm implementation has 1 layer. Cudnn implementation was implemented for 1 physical layer
// Cudnn helper restrictions:
// - all NDArrays should be the same type
// - dataFormat should be 0 or 1
// - only unidirectional (directionMode == 0) and bidirectional concat (directionMode == 3)
// - no peephole connection
// - Clipping is allowed for cudnn version >= 7.2.1
// - SeqLen array is allowed for cudnn version >= 8.0.1
// - gateActivation: sigmoid, cellActivation and outputActivation: tanh
// - NDArrays (excluding the weight arrays, as we have to transpose or permute it) should follow 'c' order and ews()==1
PLATFORM_CHECK(lstmLayer, ENGINE_CUDA) {
const auto dataFormat = INT_ARG(0); // for unidirectional: 0 = [sL, bS, nIn], 1 = [bS, sL ,nIn], 2 = [bS, nIn, sL],
// for bidirectional: 3 = [sL, 2, bS, nOut] (for ONNX)
const auto directionMode =
INT_ARG(1); // direction: 0 = fwd, 1 = bwd, 2 = bidirectional sum, 3 = bidirectional concat, 4 = bidirectional
// extra output dim (in conjunction with format dataFormat = 3)
// integer numbers corresponding to activations: 0=tanh, 1=relu, 2=sigmoid, 3=affine, 4=leaky relu, 5= thresholded
// relu, 6=scaled tanh, 7=hard sigmoid, 8=ELU, 9=softsign, 10=softplus
const auto gateAct = INT_ARG(2); // activation for input (i), forget (f) and output (o) gates
const auto cellAct = INT_ARG(3); // activation for cell state (c)
const auto outAct = INT_ARG(4); // activation for output (h)
const auto hasBiases = B_ARG(0); // indicates whether biases array is provided
const auto hasSeqLenArray = B_ARG(1); // indicates whether seqLen array is provided
const auto hasInitH = B_ARG(2); // indicates whether initial output is provided
const auto hasInitC = B_ARG(3); // indicates whether initial cell state is provided
const auto hasPH = B_ARG(4); // indicates whether peephole connections are present
const auto retFullSeq = B_ARG(5); // indicates whether to return whole time sequence h {h_0, h_1, ... , h_sL-1}
const auto retLastH = B_ARG(6); // indicates whether to return output at last time step only, in this case shape
// would be [bS, nOut] (exact shape depends on dataFormat argument)
const auto retLastC = B_ARG(7); // indicates whether to return cells state at last time step only, in this case shape
// would be [bS, nOut] (exact shape depends on dataFormat argument)
const auto cellClip = T_ARG(0); // cell clipping value, if it = 0 then do not apply clipping
const auto x = INPUT_VARIABLE(0); // input
const auto Wx = INPUT_VARIABLE(1); // input weights
const auto Wr = INPUT_VARIABLE(2); // recurrent weights
int count = 3;
const auto b = hasBiases ? INPUT_VARIABLE(count++) : nullptr; // biases
const auto hI = hasInitH ? INPUT_VARIABLE(count++) : nullptr; // initial output
const auto cI = hasInitC ? INPUT_VARIABLE(count++) : nullptr; // initial cell state
count = 0;
auto h = retFullSeq ? OUTPUT_VARIABLE(count++) : nullptr; // output
auto hL = retLastH ? OUTPUT_VARIABLE(count++) : nullptr; // output at last step
auto cL = retLastC ? OUTPUT_VARIABLE(count++) : nullptr; // cell state at last step
DataType xType = x->dataType();
DataType WxType = Wx->dataType();
DataType WrType = Wr->dataType();
Requirements req("CUDNN LSTMLAYER OP");
// cudnn related restrictions //gateAct: sigmoid, cellAct: tanh adn et cetera
// integer numbers corresponding to activations: 0=tanh, 1=relu, 2=sigmoid, 3=affine,
// 4=leaky relu, 5= thresholded relu, 6=scaled tanh, 7=hard sigmoid, 8=ELU, 9=softsign, 10=softplus
req.expectEq(makeInfoVariable(gateAct, "gate Activation"), makeInfoVariable(2, "sigmoid")) &&
req.expectEq(makeInfoVariable(cellAct, "cell Activation"), makeInfoVariable(2, "tanh")) &&
req.expectEq(makeInfoVariable(outAct, "out Activation"), makeInfoVariable(2, "tanh")) &&
req.expectFalse(makeInfoVariable(hasPH, HAVE_PEEPHOLE), EXPECTED_NOT_SUPPORTED) &&
req.expectIn(makeInfoVariable(directionMode, "directionMode"), {0, 3}) &&
req.expectIn(makeInfoVariable(dataFormat, "data Format"), {0, 1});
if (req) {
// cudnn api version related restrictions in our helpers
size_t cudnn_version = cudnnGetVersion();
// though seqlengthArray was added in earlier versions we do not handle it below 8.0.0.1
#if CUDNN_VERSION < CUDNN_NEW_RNN_API_VER
// implRestrictions = implRestrictions && !hasSeqLenArray;
req.expectFalse(makeInfoVariable(hasSeqLenArray, HAVE_SEQLENARR), EXPECTED_NOT_SUPPORTED);
#else
// implRestrictions = implRestrictions && (cudnn_version >= CUDNN_NEW_RNN_API_VER || !hasSeqLenArray);
if (cudnn_version < CUDNN_NEW_RNN_API_VER) {
req.expectFalse(makeInfoVariable(hasSeqLenArray, HAVE_SEQLENARR), EXPECTED_NOT_SUPPORTED);
}
#endif
// implRestrictions = implRestrictions && (cudnn_version >= CUDNN_CLIPPING_API_VER || cellClip==0);
if (cudnn_version < CUDNN_CLIPPING_API_VER) {
req.expectEq(makeInfoVariable(cellClip, MSG_CELL_CLIPPING), 0);
}
}
// restriction that comes either from not setting Descriptor or not handling manipulation:
// restrict0: the same types
req.expectEq(makeInfoVariable(x->ordering(), ORDERING_MSG_INPUT0), 'c') &&
req.expectEq(makeInfoVariable(WxType, TYPE_MSG_INPUT1), makeInfoVariable(xType, TYPE_MSG_INPUT0)) &&
req.expectEq(makeInfoVariable(WrType, TYPE_MSG_INPUT2), makeInfoVariable(xType, TYPE_MSG_INPUT0));
if (b)
req.expectEq(makeInfoVariable(b->dataType(), TYPE_MSG_INPUT_ "#bias"), makeInfoVariable(xType, TYPE_MSG_INPUT0));
if (hI) {
req.expectEq(makeInfoVariable(hI->dataType(), TYPE_MSG_INPUT_ "#hI"), makeInfoVariable(xType, TYPE_MSG_INPUT0)) &&
req.expectEq(makeInfoVariable(hI->ordering(), ORDERING_MSG_INPUT_ "#hI"), 'c') &&
}
if (cI) {
req.expectEq(makeInfoVariable(cI->dataType(), TYPE_MSG_INPUT_ "#cI"), makeInfoVariable(xType, TYPE_MSG_INPUT0)) &&
req.expectEq(makeInfoVariable(cI->ordering(), ORDERING_MSG_INPUT_ "#cI"), 'c') &&
}
if (h) {
req.expectEq(makeInfoVariable(h->dataType(), TYPE_MSG_OUTPUT_ "#h"), makeInfoVariable(xType, TYPE_MSG_INPUT0)) &&
req.expectEq(makeInfoVariable(h->ordering(), ORDERING_MSG_OUTPUT_ "#h"), 'c') &&
}
if (hL) {
req.expectEq(makeInfoVariable(hL->dataType(), TYPE_MSG_OUTPUT_ "#hL"), makeInfoVariable(xType, TYPE_MSG_INPUT0)) &&
req.expectEq(makeInfoVariable(hL->ordering(), ORDERING_MSG_OUTPUT_ "#hL"), 'c') &&
}
if (cL) {
req.expectEq(makeInfoVariable(cL->dataType(), TYPE_MSG_OUTPUT_ "#cL"), makeInfoVariable(xType, TYPE_MSG_INPUT0)) &&
req.expectEq(makeInfoVariable(cL->ordering(), ORDERING_MSG_OUTPUT_ "#cL"), 'c') &&
}
req.logTheSuccess();
return req;
}
} // namespace platforms
} // namespace ops
} // namespace sd
@@ -0,0 +1,157 @@
/* ******************************************************************************
*
*
* 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)
//
#include <ops/declarable/helpers/convolutions.h>
#include "cudnnUtils.h"
namespace sd {
namespace ops {
namespace platforms {
//////////////////////////////////////////////////////////////////////////
PLATFORM_IMPL(maxpool2d, ENGINE_CUDA) {
auto input = INPUT_VARIABLE(0);
auto output = OUTPUT_VARIABLE(0);
// 0,1 - kernel Height/Width; 2,3 - stride Height/Width; 4,5 - pad Height/Width; 6,7 - dilation Height/Width; 8 -
// paddingModee;
const LongType kH = INT_ARG(0);
const LongType kW = INT_ARG(1);
const LongType sH = INT_ARG(2);
const LongType sW = INT_ARG(3);
LongType pH = INT_ARG(4);
LongType pW = INT_ARG(5);
const LongType dH = INT_ARG(6);
const LongType dW = INT_ARG(7);
const auto paddingMode = static_cast<bool>(INT_ARG(8));
const int isNCHW = block.getIArguments()->size() > 10 ? !INT_ARG(10) : 1; // INT_ARG(10): 0-NCHW, 1-NHWC
REQUIRE_TRUE(input->rankOf() == 4, 0, "MAXPOOL2D CUDNN op: input should have rank of 4, but got %i instead",
input->rankOf());
REQUIRE_TRUE(dH != 0 && dW != 0, 0, "MAXPOOL2D CUDNN op: dilation must not be zero, but got instead {%i, %i}", dH,
dW);
LongType oH = 0;
LongType oW = 0;
const LongType iH = static_cast<LongType>(isNCHW ? input->sizeAt(2) : input->sizeAt(1));
const LongType iW = static_cast<LongType>(isNCHW ? input->sizeAt(3) : input->sizeAt(2));
ConvolutionUtils::calcOutSizePool2D(oH, oW, kH, kW, sH, sW, pH, pW, dH, dW, iH, iW, paddingMode);
if (paddingMode) ConvolutionUtils::calcPadding2D(pH, pW, oH, oW, iH, iW, kH, kW, sH, sW, dH, dW);
pooling2dCUDNN(block.launchContext(), input, output, kH, kW, sH, sW, pH, pW, dH, dW, isNCHW, CUDNN_POOLING_MAX);
return Status::OK;
}
//////////////////////////////////////////////////////////////////////////
PLATFORM_CHECK(maxpool2d, ENGINE_CUDA) {
auto input = INPUT_VARIABLE(0);
auto output = OUTPUT_VARIABLE(0);
Requirements req("CUDNN MAXPOOL2d OP");
req.expectEq(makeInfoVariable(input->dataType(), TYPE_MSG_INPUT),
makeInfoVariable(output->dataType(), TYPE_MSG_OUTPUT)) &&
req.expectIn(makeInfoVariable(input->dataType(), TYPE_MSG_INPUT),
{INT32, HALF, FLOAT32, DOUBLE});
req.logTheSuccess();
return req;
}
//////////////////////////////////////////////////////////////////////////
PLATFORM_IMPL(maxpool2d_bp, ENGINE_CUDA) {
auto input = INPUT_VARIABLE(0); // [bS, iH, iW, iC] (NHWC) or [bS, iC, iH, iW] (NCHW)
auto gradO = INPUT_VARIABLE(1); // [bS, oH, oW, oC] (NHWC) or [bS, oC, oH, oW] (NCHW), epsilon_next
auto gradI = OUTPUT_VARIABLE(0); // [bS, iH, iW, iC] (NHWC) or [bS, iC, iH, iW] (NCHW), epsilon
const LongType kH = INT_ARG(0); // filter(kernel) height
const LongType kW = INT_ARG(1); // filter(kernel) width
const LongType sH = INT_ARG(2); // strides height
const LongType sW = INT_ARG(3); // strides width
LongType pH = INT_ARG(4); // paddings height
LongType pW = INT_ARG(5); // paddings width
const LongType dH = INT_ARG(6); // dilations height
const LongType dW = INT_ARG(7); // dilations width
const auto paddingMode = INT_ARG(8); // 0-VALID, 1-SAME
const auto isNCHW = block.getIArguments()->size() > 10 ? !INT_ARG(10) : 1; // INT_ARG(10): 0-NCHW, 1-NHWC
REQUIRE_TRUE(input->rankOf() == 4, 0, "MAXPOOL2D_BP CUDNN op: input should have rank of 4, but got %i instead",
input->rankOf());
REQUIRE_TRUE(dH != 0 && dW != 0, 0, "MAXPOOL2D_BP CUDNN op: dilation must not be zero, but got instead {%i, %i}", dH,
dW);
LongType bS, iC, iH, iW, oC, oH,
oW; // batch size, input channels, input height/width, output channels, output height/width;
LongType indIOioC, indIiH, indWoC, indWiC, indWkH, indOoH; // corresponding indexes
ConvolutionUtils::getSizesAndIndexesConv2d(isNCHW, 0, *input, *gradO, bS, iC, iH, iW, oC, oH, oW, indIOioC, indIiH,
indWiC, indWoC, indWkH, indOoH);
std::vector<LongType> expectedGradOShape =
ShapeUtils::composeShapeUsingDimsAndIdx({bS, iC, oH, oW, 0, indIOioC, indIiH, indIiH + 1});
std::vector<LongType> expectedGradIShape =
ShapeUtils::composeShapeUsingDimsAndIdx({bS, iC, iH, iW, 0, indIOioC, indIiH, indIiH + 1});
REQUIRE_TRUE(gradO->isSameShape(expectedGradOShape), 0,
"MAXPOOL2D_BP CUDNN op: wrong shape of output's gradients array (next epsilon), expected is %s, but got "
"%s instead !",
ShapeUtils::shapeAsString(expectedGradOShape).c_str(), ShapeUtils::shapeAsString(gradO).c_str());
REQUIRE_TRUE(
gradI->isSameShape(expectedGradIShape), 0,
"MAXPOOL2D_BP CUDNN op: wrong shape of input's gradients array (epsilon), expected is %s, but got %s instead !",
ShapeUtils::shapeAsString(expectedGradIShape).c_str(), ShapeUtils::shapeAsString(gradI).c_str());
if (paddingMode) // SAME
ConvolutionUtils::calcPadding2D(pH, pW, oH, oW, iH, iW, kH, kW, sH, sW, dH, dW);
pooling2dBpCUDNN(block.launchContext(), input, gradO, gradI, kH, kW, sH, sW, pH, pW, dH, dW, isNCHW,
CUDNN_POOLING_MAX);
return Status::OK;
}
PLATFORM_CHECK(maxpool2d_bp, ENGINE_CUDA) {
auto input = INPUT_VARIABLE(0); // [bS, iH, iW, iC] (NHWC) or [bS, iC, iH, iW] (NCHW)
auto gradO = INPUT_VARIABLE(1); // [bS, oH, oW, oC] (NHWC) or [bS, oC, oH, oW] (NCHW), epsilon_next
auto gradI = OUTPUT_VARIABLE(0); // [bS, iH, iW, iC] (NHWC) or [bS, iC, iH, iW] (NCHW), epsilon
Requirements req("CUDNN MAXPOOL2d_BP OP");
req.expectEq(makeInfoVariable(input->ordering(), ORDERING_MSG_INPUT), 'c') &&
req.expectEq(makeInfoVariable(input->dataType(), TYPE_MSG_INPUT0),
makeInfoVariable(gradO->dataType(), TYPE_MSG_INPUT1)) &&
req.expectEq(makeInfoVariable(input->dataType(), TYPE_MSG_INPUT),
makeInfoVariable(gradI->dataType(), TYPE_MSG_OUTPUT)) &&
req.expectIn(makeInfoVariable(input->dataType(), TYPE_MSG_INPUT),
{INT32, HALF, FLOAT32, DOUBLE}) &&
req.expect(
makeShapeInfoVariable(input, SHAPE_MSG_INPUT0), makeShapeInfoVariable(gradI, SHAPE_MSG_OUTPUT),
[](const decltype(input)& l, const decltype(gradI)& r) {
return shape::haveSameShapeAndStrides(l->shapeInfo(), r->shapeInfo());
},
EXPECTED_EQ_MSG);
req.logTheSuccess();
return req;
}
} // namespace platforms
} // namespace ops
} // namespace sd
@@ -0,0 +1,171 @@
/* ******************************************************************************
*
*
* 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)
//
#include <ops/declarable/helpers/convolutions.h>
#include "cudnnUtils.h"
namespace sd {
namespace ops {
namespace platforms {
//////////////////////////////////////////////////////////////////////////
PLATFORM_IMPL(maxpool3dnew, ENGINE_CUDA) {
auto input = INPUT_VARIABLE(0); // [bS, iD, iH, iW, iC] (NDHWC) or [bS, iC, iD, iH, iW] (NCDHW)
auto output = OUTPUT_VARIABLE(0); // [bS, oD, oH, oW, iC] (NDHWC) or [bS, iC, oD, oH, oW] (NCDHW)
LongType kD = INT_ARG(0); // filter(kernel) depth
LongType kH = INT_ARG(1); // filter(kernel) height
LongType kW = INT_ARG(2); // filter(kernel) width
LongType sD = INT_ARG(3); // strides depth
LongType sH = INT_ARG(4); // strides height
LongType sW = INT_ARG(5); // strides width
LongType pD = INT_ARG(6); // paddings depth
LongType pH = INT_ARG(7); // paddings height
LongType pW = INT_ARG(8); // paddings width
LongType dD = INT_ARG(9); // dilations depth
LongType dH = INT_ARG(10); // dilations height
LongType dW = INT_ARG(11); // dilations width
int paddingMode = INT_ARG(12); // 1-SAME, 0-VALID
// int extraParam0 = INT_ARG(13);
int isNCDHW = block.getIArguments()->size() > 14 ? !INT_ARG(14) : 1; // 0-NCDHW, 1-NDHWC
REQUIRE_TRUE(input->rankOf() == 5, 0,
"MAXPOOL3DNEW CUDNN OP: rank of input array must be equal to 5, but got %i instead !", input->rankOf());
REQUIRE_TRUE(dD != 0 && dH != 0 && dW != 0, 0,
"MAXPOOL3DNEW CUDNN OP: dilation must not be zero, but got instead {%i, %i, %i}", dD, dH, dW);
LongType bS, iC, iD, iH, iW, oC, oD, oH,
oW; // batch size, input channels, input depth/height/width, output channels, output depth/height/width;
LongType indIOioC, indIOioD, indWoC, indWiC, indWkD; // corresponding indexes
ConvolutionUtils::getSizesAndIndexesConv3d(isNCDHW, 0, *input, *output, bS, iC, iD, iH, iW, oC, oD, oH, oW, indIOioC,
indIOioD, indWiC, indWoC, indWkD);
std::vector<LongType> expectedOutputShape =
ShapeUtils::composeShapeUsingDimsAndIdx({bS, iC, oD, oH, oW, 0, indIOioC, indIOioD, indIOioD + 1, indIOioD + 2});
REQUIRE_TRUE(output->isSameShape(expectedOutputShape), 0,
"MAXPOOL3DNEW CUDNN OP: wrong shape of output array, expected is %s, but got %s instead !",
ShapeUtils::shapeAsString(expectedOutputShape).c_str(), ShapeUtils::shapeAsString(output).c_str());
if (paddingMode) // SAME
ConvolutionUtils::calcPadding3D(pD, pH, pW, oD, oH, oW, iD, iH, iW, kD, kH, kW, sD, sH, sW, dD, dH, dW);
pooling3dCUDNN(block.launchContext(), input, output, kD, kH, kW, sD, sH, sW, pD, pH, pW, dD, dH, dW, isNCDHW,
CUDNN_POOLING_MAX);
return Status::OK;
}
//////////////////////////////////////////////////////////////////////////
PLATFORM_CHECK(maxpool3dnew, ENGINE_CUDA) {
auto input = INPUT_VARIABLE(0);
auto output = OUTPUT_VARIABLE(0);
Requirements req("CUDNN MAXPOOL3d OP");
req.expectEq(makeInfoVariable(input->dataType(), TYPE_MSG_INPUT),
makeInfoVariable(output->dataType(), TYPE_MSG_OUTPUT)) &&
req.expectIn(makeInfoVariable(input->dataType(), TYPE_MSG_INPUT),
{INT32, HALF, FLOAT32, DOUBLE});
req.logTheSuccess();
return req;
}
//////////////////////////////////////////////////////////////////////////
PLATFORM_IMPL(maxpool3dnew_bp, ENGINE_CUDA) {
auto input = INPUT_VARIABLE(0); // [bS, iD, iH, iW, iC] (NDHWC) or [bS, iC, iD, iH, iW] (NCDHW)
auto gradO = INPUT_VARIABLE(1); // [bS, oD, oH, oW, oC] (NDHWC) or [bS, oC, oD, oH, oW] (NCDHW), epsilon_next
auto gradI = OUTPUT_VARIABLE(0); // [bS, iD, iH, iW, iC] (NDHWC) or [bS, iC, iD, iH, iW] (NCDHW), epsilon
const LongType kD = INT_ARG(0); // filter(kernel) depth
const LongType kH = INT_ARG(1); // filter(kernel) height
const LongType kW = INT_ARG(2); // filter(kernel) width
const LongType sD = INT_ARG(3); // strides depth
const LongType sH = INT_ARG(4); // strides height
const LongType sW = INT_ARG(5); // strides width
LongType pD = INT_ARG(6); // paddings depth
LongType pH = INT_ARG(7); // paddings height
LongType pW = INT_ARG(8); // paddings width
const LongType dD = INT_ARG(9); // dilations depth
const LongType dH = INT_ARG(10); // dilations height
const LongType dW = INT_ARG(11); // dilations width
const int isSameMode = INT_ARG(12); // 1-SAME, 0-VALID
// const int extraParam0 = INT_ARG(13); // define what divisor to use while
// averaging
const int isNCDHW = block.getIArguments()->size() > 14 ? !INT_ARG(14) : 1; // 0-NCDHW, 1-NDHWC
REQUIRE_TRUE(input->rankOf() == 5, 0, "MAXPOOL3DNEW_BP CUDNN OP: input should have rank of 5, but got %i instead",
input->rankOf());
REQUIRE_TRUE(dD != 0 && dH != 0 && dW != 0, 0,
"MAXPOOL3DNEW_BP CUDNN OP: dilation must not be zero, but got instead {%i, %i, %i}", dD, dH, dW);
LongType bS, iC, iD, iH, iW, oC, oD, oH,
oW; // batch size, input channels, input depth/height/width, output channels, output depth/height/width;
LongType indIOioC, indIOioD, indWoC, indWiC, indWkD; // corresponding indexes
ConvolutionUtils::getSizesAndIndexesConv3d(isNCDHW, 0, *input, *gradO, bS, iC, iD, iH, iW, oC, oD, oH, oW, indIOioC,
indIOioD, indWiC, indWoC, indWkD);
std::vector<LongType> expectedGradOShape =
ShapeUtils::composeShapeUsingDimsAndIdx({bS, iC, oD, oH, oW, 0, indIOioC, indIOioD, indIOioD + 1, indIOioD + 2});
std::vector<LongType> expectedGradIShape =
ShapeUtils::composeShapeUsingDimsAndIdx({bS, iC, iD, iH, iW, 0, indIOioC, indIOioD, indIOioD + 1, indIOioD + 2});
REQUIRE_TRUE(gradO->isSameShape(expectedGradOShape), 0,
"MAXPOOL3DNEW_BP CUDNN: wrong shape of output's gradients array (next epsilon), expected is %s, but got "
"%s instead !",
ShapeUtils::shapeAsString(expectedGradOShape).c_str(), ShapeUtils::shapeAsString(gradO).c_str());
REQUIRE_TRUE(
gradI->isSameShape(expectedGradIShape), 0,
"MAXPOOL3DNEW_BP CUDNN: wrong shape of input's gradients array (epsilon), expected is %s, but got %s instead !",
ShapeUtils::shapeAsString(expectedGradIShape).c_str(), ShapeUtils::shapeAsString(gradI).c_str());
if (isSameMode) // SAME
ConvolutionUtils::calcPadding3D(pD, pH, pW, oD, oH, oW, iD, iH, iW, kD, kH, kW, sD, sH, sW, dD, dH, dW);
pooling3dBpCUDNN(block.launchContext(), input, gradO, gradI, kD, kH, kW, sD, sH, sW, pD, pH, pW, dD, dH, dW, isNCDHW,
CUDNN_POOLING_MAX);
return Status::OK;
}
PLATFORM_CHECK(maxpool3dnew_bp, ENGINE_CUDA) {
auto input = INPUT_VARIABLE(0); // [bS, iD, iH, iW, iC] (NDHWC) or [bS, iC, iD, iH, iW] (NCDHW)
auto gradO = INPUT_VARIABLE(1); // [bS, oD, oH, oW, oC] (NDHWC) or [bS, oC, oD, oH, oW] (NCDHW), epsilon_next
auto gradI = OUTPUT_VARIABLE(0); // [bS, iD, iH, iW, iC] (NDHWC) or [bS, iC, iD, iH, iW] (NCDHW), epsilon
Requirements req("CUDNN MAXPOOL3d_BP OP");
req.expectEq(makeInfoVariable(input->dataType(), TYPE_MSG_INPUT0),
makeInfoVariable(gradO->dataType(), TYPE_MSG_INPUT1)) &&
req.expectEq(makeInfoVariable(input->dataType(), TYPE_MSG_INPUT),
makeInfoVariable(gradI->dataType(), TYPE_MSG_OUTPUT)) &&
req.expectIn(makeInfoVariable(input->dataType(), TYPE_MSG_INPUT),
{INT32, HALF, FLOAT32, DOUBLE}) &&
req.expect(
makeShapeInfoVariable(input, SHAPE_MSG_INPUT0), makeShapeInfoVariable(gradI, SHAPE_MSG_OUTPUT),
[](const decltype(input)& l, const decltype(gradI)& r) {
return shape::haveSameShapeAndStrides(l->shapeInfo(), r->shapeInfo());
},
EXPECTED_EQ_MSG);
req.logTheSuccess();
return req;
}
} // namespace platforms
} // namespace ops
} // namespace sd