chore: import upstream snapshot with attribution
This commit is contained in:
@@ -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
|
||||
Reference in New Issue
Block a user