chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,152 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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 saudet
|
||||
// @author raver119@gmail.com
|
||||
// @author Yurii Shyrma (iuriish@yahoo.com)
|
||||
//
|
||||
#include <helpers/MKLDNNStream.h>
|
||||
#include <ops/declarable/OpRegistrator.h>
|
||||
#include <ops/declarable/PlatformHelper.h>
|
||||
#include <ops/declarable/helpers/convolutions.h>
|
||||
#include <system/platform_boilerplate.h>
|
||||
|
||||
#include "mkldnnUtils.h"
|
||||
|
||||
using namespace dnnl;
|
||||
using namespace samediff;
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace platforms {
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
PLATFORM_IMPL(avgpool2d, ENGINE_CPU) {
|
||||
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 sd::LongType kH = INT_ARG(0);
|
||||
const sd::LongType kW = INT_ARG(1);
|
||||
const sd::LongType sH = INT_ARG(2);
|
||||
const sd::LongType sW = INT_ARG(3);
|
||||
sd::LongType pH = INT_ARG(4);
|
||||
sd::LongType pW = INT_ARG(5);
|
||||
const sd::LongType dH = INT_ARG(6);
|
||||
const sd::LongType dW = INT_ARG(7);
|
||||
const auto paddingMode = 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 MKLDNN op: input should have rank of 4, but got %i instead",
|
||||
input->rankOf());
|
||||
REQUIRE_TRUE(dH != 0 && dW != 0, 0, "AVGPOOL2D MKLDNN op: dilation must not be zero, but got instead {%i, %i}", dH,
|
||||
dW);
|
||||
|
||||
sd::LongType bS, iC, iH, iW, oC, oH,
|
||||
oW; // batch size, input channels, input height/width, output channels, output height/width;
|
||||
sd::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);
|
||||
|
||||
if (paddingMode) ConvolutionUtils::calcPadding2D(pH, pW, oH, oW, iH, iW, kH, kW, sH, sW, dH, dW);
|
||||
|
||||
auto mode = (extraParam0 == 0) ? algorithm::pooling_avg_exclude_padding : algorithm::pooling_avg_include_padding;
|
||||
|
||||
onednnUtils::poolingONEDNN(input, output, 0, kH, kW, 0, sH, sW, 0, pH, pW, isNCHW, mode);
|
||||
|
||||
return sd::Status::OK;
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
PLATFORM_CHECK(avgpool2d, ENGINE_CPU) {
|
||||
auto input = INPUT_VARIABLE(0);
|
||||
auto output = OUTPUT_VARIABLE(0);
|
||||
Requirements req("ONEDNN AVGPOOL2d OP");
|
||||
req.expectTrue(block.isUseONEDNN(), IS_USE_ONEDNN_MSG) &&
|
||||
req.expectTrue(sd::ONEDNNStream::isSupported({input, output}), ONEDNN_STREAM_NOT_SUPPORTED);
|
||||
if (req) onednnUtils::checkPoolingONEDNN(req, block, input, output);
|
||||
req.logTheSuccess();
|
||||
return req;
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
PLATFORM_IMPL(avgpool2d_bp, ENGINE_CPU) {
|
||||
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
|
||||
|
||||
sd::LongType kH = INT_ARG(0); // filter(kernel) height
|
||||
sd::LongType kW = INT_ARG(1); // filter(kernel) width
|
||||
sd::LongType sH = INT_ARG(2); // strides height
|
||||
sd::LongType sW = INT_ARG(3); // strides width
|
||||
sd::LongType pH = INT_ARG(4); // paddings height
|
||||
sd::LongType pW = INT_ARG(5); // paddings width
|
||||
sd::LongType dH = INT_ARG(6); // dilations height
|
||||
sd::LongType dW = INT_ARG(7); // dilations width
|
||||
int paddingMode = INT_ARG(8); // 0-VALID, 1-SAME
|
||||
int extraParam0 = INT_ARG(9);
|
||||
int isNCHW = block.getIArguments()->size() > 10 ? !INT_ARG(10) : 1; // INT_ARG(10): 0-NCHW, 1-NHWC
|
||||
|
||||
REQUIRE_TRUE(input->rankOf() == 4, 0, "AVGPOOL2D_BP MKLDNN op: input should have rank of 4, but got %i instead",
|
||||
input->rankOf());
|
||||
REQUIRE_TRUE(dH != 0 && dW != 0, 0, "AVGPOOL2D_BP MKLDNN op: dilation must not be zero, but got instead {%i, %i}", dH,
|
||||
dW);
|
||||
|
||||
sd::LongType bS, iC, iH, iW, oC, oH,
|
||||
oW; // batch size, input channels, input height/width, output channels, output height/width;
|
||||
sd::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<sd::LongType> expectedGradOShape =
|
||||
ShapeUtils::composeShapeUsingDimsAndIdx({bS, iC, oH, oW, 0, indIOioC, indIiH, indIiH + 1});
|
||||
REQUIRE_TRUE(gradO->isSameShape(expectedGradOShape), 0,
|
||||
"AVGPOOL2D_BP MKLDNN 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());
|
||||
|
||||
if (paddingMode) // SAME
|
||||
ConvolutionUtils::calcPadding2D(pH, pW, oH, oW, iH, iW, kH, kW, sH, sW, dH, dW);
|
||||
|
||||
auto mode = (extraParam0 == 0) ? algorithm::pooling_avg_exclude_padding : algorithm::pooling_avg_include_padding;
|
||||
|
||||
onednnUtils::poolingBpONEDNN(input, gradO, gradI, 0, kH, kW, 0, sH, sW, 0, pH, pW, isNCHW, mode);
|
||||
|
||||
return sd::Status::OK;
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
PLATFORM_CHECK(avgpool2d_bp, ENGINE_CPU) {
|
||||
auto input = INPUT_VARIABLE(0);
|
||||
auto gradO = INPUT_VARIABLE(1);
|
||||
auto output = OUTPUT_VARIABLE(0);
|
||||
Requirements req("ONEDNN AVGPOOL2d_BP OP");
|
||||
req.expectTrue(block.isUseONEDNN(), IS_USE_ONEDNN_MSG) &&
|
||||
req.expectTrue(sd::ONEDNNStream::isSupported({input, output}), ONEDNN_STREAM_NOT_SUPPORTED);
|
||||
if (req) onednnUtils::checkPoolingONEDNN(req, block,input, gradO);
|
||||
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 saudet
|
||||
// @author raver119@gmail.com
|
||||
// @author Yurii Shyrma (iuriish@yahoo.com)
|
||||
//
|
||||
#include <helpers/MKLDNNStream.h>
|
||||
#include <ops/declarable/OpRegistrator.h>
|
||||
#include <ops/declarable/PlatformHelper.h>
|
||||
#include <ops/declarable/helpers/convolutions.h>
|
||||
#include <system/platform_boilerplate.h>
|
||||
|
||||
#include "mkldnnUtils.h"
|
||||
|
||||
using namespace dnnl;
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace platforms {
|
||||
|
||||
//////////////////////////////////////////////////////////////////////
|
||||
PLATFORM_IMPL(avgpool3dnew, ENGINE_CPU) {
|
||||
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)
|
||||
|
||||
sd::LongType kD = INT_ARG(0); // filter(kernel) depth
|
||||
sd::LongType kH = INT_ARG(1); // filter(kernel) height
|
||||
sd::LongType kW = INT_ARG(2); // filter(kernel) width
|
||||
sd::LongType sD = INT_ARG(3); // strides depth
|
||||
sd::LongType sH = INT_ARG(4); // strides height
|
||||
sd::LongType sW = INT_ARG(5); // strides width
|
||||
sd::LongType pD = INT_ARG(6); // paddings depth
|
||||
sd::LongType pH = INT_ARG(7); // paddings height
|
||||
sd::LongType pW = INT_ARG(8); // paddings width
|
||||
sd::LongType dD = INT_ARG(9); // dilations depth
|
||||
sd::LongType dH = INT_ARG(10); // dilations height
|
||||
sd::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 MKLDNN 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 MKLDNN OP: dilation must not be zero, but got instead {%i, %i, %i}", dD, dH, dW);
|
||||
|
||||
sd::LongType bS, iC, iD, iH, iW, oC, oD, oH,
|
||||
oW; // batch size, input channels, input depth/height/width, output channels, output depth/height/width;
|
||||
sd::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);
|
||||
|
||||
if (paddingMode) // SAME
|
||||
ConvolutionUtils::calcPadding3D(pD, pH, pW, oD, oH, oW, iD, iH, iW, kD, kH, kW, sD, sH, sW, dD, dH, dW);
|
||||
|
||||
auto mode = (extraParam0 == 0) ? algorithm::pooling_avg_exclude_padding : algorithm::pooling_avg_include_padding;
|
||||
|
||||
onednnUtils::poolingONEDNN(input, output, kD, kH, kW, sD, sH, sW, pD, pH, pW, isNCDHW, mode);
|
||||
|
||||
return sd::Status::OK;
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////
|
||||
PLATFORM_CHECK(avgpool3dnew, ENGINE_CPU) {
|
||||
auto input = INPUT_VARIABLE(0);
|
||||
auto output = OUTPUT_VARIABLE(0);
|
||||
Requirements req("ONEDNN AVGPOOL3d OP");
|
||||
req.expectTrue(block.isUseONEDNN(), IS_USE_ONEDNN_MSG) &&
|
||||
req.expectTrue(sd::ONEDNNStream::isSupported({input, output}), ONEDNN_STREAM_NOT_SUPPORTED);
|
||||
if (req) onednnUtils::checkPoolingONEDNN(req, block, input, output);
|
||||
req.logTheSuccess();
|
||||
return req;
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////
|
||||
PLATFORM_IMPL(avgpool3dnew_bp, ENGINE_CPU) {
|
||||
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 sd::LongType kD = INT_ARG(0); // filter(kernel) depth
|
||||
const sd::LongType kH = INT_ARG(1); // filter(kernel) height
|
||||
const sd::LongType kW = INT_ARG(2); // filter(kernel) width
|
||||
const sd::LongType sD = INT_ARG(3); // strides depth
|
||||
const sd::LongType sH = INT_ARG(4); // strides height
|
||||
const sd::LongType sW = INT_ARG(5); // strides width
|
||||
sd::LongType pD = INT_ARG(6); // paddings depth
|
||||
sd::LongType pH = INT_ARG(7); // paddings height
|
||||
sd::LongType pW = INT_ARG(8); // paddings width
|
||||
const sd::LongType dD = INT_ARG(9); // dilations depth
|
||||
const sd::LongType dH = INT_ARG(10); // dilations height
|
||||
const sd::LongType dW = INT_ARG(11); // dilations width
|
||||
const int paddingMode = 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 MKLDNN op: input should have rank of 5, but got %i instead",
|
||||
input->rankOf());
|
||||
REQUIRE_TRUE(dD != 0 && dH != 0 && dW != 0, 0,
|
||||
"AVGPOOL3DNEW_BP MKLDNN op: dilation must not be zero, but got instead {%i, %i, %i}", dD, dH, dW);
|
||||
|
||||
sd::LongType bS, iC, iD, iH, iW, oC, oD, oH,
|
||||
oW; // batch size, input channels, input depth/height/width, output channels, output depth/height/width;
|
||||
sd::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<sd::LongType> expectedGradOShape =
|
||||
ShapeUtils::composeShapeUsingDimsAndIdx({bS, iC, oD, oH, oW, 0, indIOioC, indIOioD, indIOioD + 1, indIOioD + 2});
|
||||
REQUIRE_TRUE(gradO->isSameShape(expectedGradOShape), 0,
|
||||
"AVGPOOL3DNEW_BP MKLDNN 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());
|
||||
|
||||
if (paddingMode) // SAME
|
||||
ConvolutionUtils::calcPadding3D(pD, pH, pW, oD, oH, oW, iD, iH, iW, kD, kH, kW, sD, sH, sW, dD, dH, dW);
|
||||
|
||||
auto mode = (extraParam0 == 0) ? algorithm::pooling_avg_exclude_padding : algorithm::pooling_avg_include_padding;
|
||||
|
||||
onednnUtils::poolingBpONEDNN(input, gradO, gradI, kD, kH, kW, sD, sH, sW, pD, pH, pW, isNCDHW, mode);
|
||||
|
||||
return sd::Status::OK;
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////
|
||||
PLATFORM_CHECK(avgpool3dnew_bp, ENGINE_CPU) {
|
||||
auto input = INPUT_VARIABLE(0);
|
||||
auto gradO = INPUT_VARIABLE(1);
|
||||
auto output = OUTPUT_VARIABLE(0);
|
||||
Requirements req("ONEDNN AVGPOOL3d_BP OP");
|
||||
req.expectTrue(block.isUseONEDNN(), IS_USE_ONEDNN_MSG) &&
|
||||
req.expectTrue(sd::ONEDNNStream::isSupported({input, output}), ONEDNN_STREAM_NOT_SUPPORTED);
|
||||
if (req) onednnUtils::checkPoolingONEDNN(req, block, input, gradO);
|
||||
req.logTheSuccess();
|
||||
return req;
|
||||
}
|
||||
|
||||
} // namespace platforms
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,639 @@
|
||||
/*
|
||||
* ******************************************************************************
|
||||
* *
|
||||
* *
|
||||
* * 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 saudet
|
||||
// @author raver119@gmail.com
|
||||
// @author Yurii Shyrma (iuriish@yahoo.com)
|
||||
//
|
||||
#include <array/NDArrayFactory.h>
|
||||
#include <helpers/MKLDNNStream.h>
|
||||
#include <ops/declarable/OpRegistrator.h>
|
||||
#include <ops/declarable/PlatformHelper.h>
|
||||
#include <ops/declarable/helpers/convolutions.h>
|
||||
#include <system/platform_boilerplate.h>
|
||||
|
||||
#include "mkldnnUtils.h"
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace platforms {
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
static void batchnormMKLDNN(NDArray* x, NDArray* mean, NDArray* variance, NDArray* weights,
|
||||
NDArray* z, const float epsilon, const bool isNCHW) {
|
||||
// unfortunately mkl dnn doesn't support any format (dnnl::memory::format_tag::any) for x
|
||||
|
||||
// x -> 2D:nc, 4D:nchw/nhwc, 5D:ncdhw/ndhwc
|
||||
// mean -> 1D [c]
|
||||
// variance -> 1D [c]
|
||||
// weights 2D [2, c], weights({0,1, 0,0}) contains gamma and weights({1,2, 0,0}) contains beta
|
||||
// z(output) - same shape as x
|
||||
|
||||
const int xRank = x->rankOf();
|
||||
|
||||
// input type
|
||||
dnnl::memory::data_type type = dnnl::memory::data_type::f32;
|
||||
|
||||
// indicate whether gamma or/and beta are given
|
||||
auto flags =
|
||||
dnnl::normalization_flags::use_global_stats; // don't calculate the mean and variance for each mini-batch
|
||||
if (weights != nullptr) flags |= dnnl::normalization_flags::use_scale_shift;
|
||||
|
||||
dnnl::memory::dims dims;
|
||||
dnnl::memory::format_tag format;
|
||||
|
||||
const sd::LongType indHW = isNCHW ? 2 : 1;
|
||||
const sd::LongType bS = x->sizeAt(0);
|
||||
const sd::LongType iC = isNCHW ? x->sizeAt(1) : x->sizeAt(-1);
|
||||
|
||||
int iD, iH, iW;
|
||||
|
||||
if (xRank == 2) {
|
||||
dims = {bS, iC};
|
||||
format = dnnl::memory::format_tag::nc;
|
||||
} else if (xRank == 4) {
|
||||
iH = x->sizeAt(indHW);
|
||||
iW = x->sizeAt(indHW + 1);
|
||||
dims = {bS, iC, iH, iW};
|
||||
format = isNCHW ? dnnl::memory::format_tag::nchw : dnnl::memory::format_tag::nhwc;
|
||||
} else { // xRank = 5
|
||||
iD = x->sizeAt(indHW);
|
||||
iH = x->sizeAt(indHW + 1);
|
||||
iW = x->sizeAt(indHW + 2);
|
||||
dims = {bS, iC, iD, iH, iW};
|
||||
format = isNCHW ? dnnl::memory::format_tag::ncdhw : dnnl::memory::format_tag::ndhwc;
|
||||
}
|
||||
|
||||
// memory descriptors for arrays
|
||||
|
||||
// x
|
||||
dnnl::memory::desc x_mkl_md = dnnl::memory::desc(dims, type, format);
|
||||
dnnl::memory::desc x_user_md = dnnl::memory::desc(dims, type, format);
|
||||
onednnUtils::setBlockStrides(*x, x_user_md);
|
||||
|
||||
// z, output
|
||||
dnnl::memory::desc z_mkl_md = dnnl::memory::desc(dims, type, dnnl::memory::format_tag::any);
|
||||
dnnl::memory::desc z_user_md = dnnl::memory::desc(dims, type, format);
|
||||
onednnUtils::setBlockStrides(*z, z_user_md);
|
||||
|
||||
auto engine = onednnUtils::getEngine(LaunchContext::defaultContext()->engine());
|
||||
|
||||
// batchnorm forward description
|
||||
dnnl::batch_normalization_forward::desc op_ff_desc(dnnl::prop_kind::forward_inference, x_mkl_md, epsilon, flags);
|
||||
dnnl::batch_normalization_forward::primitive_desc op_ff_prim_desc(op_ff_desc, engine);
|
||||
|
||||
// arguments (memory buffers) necessary for calculations
|
||||
std::unordered_map<int, dnnl::memory> args;
|
||||
|
||||
dnnl::stream stream(engine);
|
||||
|
||||
// provide memory and check whether reorder is required
|
||||
|
||||
// x
|
||||
onednnUtils::loadDataToMklStream(*x, engine, stream, x_user_md, op_ff_prim_desc.src_desc(), args[DNNL_ARG_SRC]);
|
||||
|
||||
// z
|
||||
auto z_user_mem =
|
||||
onednnUtils::loadDataToMklStream(*z, engine, stream, z_user_md, op_ff_prim_desc.dst_desc(), args[DNNL_ARG_DST]);
|
||||
|
||||
// mean
|
||||
auto mean_mkl_mem = dnnl::memory(op_ff_prim_desc.mean_desc(), engine, const_cast<void*>(mean->buffer()));
|
||||
args[DNNL_ARG_MEAN] = mean_mkl_mem;
|
||||
|
||||
// variance
|
||||
auto var_mkl_mem = dnnl::memory(op_ff_prim_desc.variance_desc(), engine, const_cast<void*>(variance->buffer()));
|
||||
args[DNNL_ARG_VARIANCE] = var_mkl_mem;
|
||||
|
||||
// gamma and beta (and their gradients) if they are present
|
||||
if (weights != nullptr) {
|
||||
auto w_mkl_mem = dnnl::memory(op_ff_prim_desc.weights_desc(), engine, const_cast<void*>(weights->buffer()));
|
||||
args[DNNL_ARG_WEIGHTS] = w_mkl_mem;
|
||||
}
|
||||
|
||||
// run calculations
|
||||
dnnl::batch_normalization_forward(op_ff_prim_desc).execute(stream, args);
|
||||
|
||||
// reorder outputs if necessary
|
||||
if (op_ff_prim_desc.dst_desc() != z_user_mem.get_desc())
|
||||
dnnl::reorder(args[DNNL_ARG_DST], z_user_mem).execute(stream, args[DNNL_ARG_DST], z_user_mem);
|
||||
|
||||
stream.wait();
|
||||
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
static void batchnormBpMKLDNN(NDArray* x, NDArray* mean, NDArray* variance, NDArray* dLdO,
|
||||
NDArray* weights, NDArray* dLdI, NDArray* dLdW, const float epsilon,
|
||||
const bool isNCHW) {
|
||||
// unfortunately mkl dnn doesn't support any format (dnnl::memory::format_tag::any) for x
|
||||
|
||||
// x -> 2D:nc, 4D:nchw/nhwc, 5D:ncdhw/ndhwc
|
||||
// mean -> 1D [c]
|
||||
// variance -> 1D [c]
|
||||
// dLdO - same shape as x
|
||||
// weights 2D [2, c], weights({0,1, 0,0}) contains gamma and weights({1,2, 0,0}) contains beta
|
||||
// dLdI - same shape as x
|
||||
// dLdW - same shape as weights, dLdW({0,1, 0,0}) contains grad_gamma and dLdW({1,2, 0,0}) contains grad_beta
|
||||
|
||||
const sd::LongType xRank = x->rankOf();
|
||||
|
||||
// input type
|
||||
dnnl::memory::data_type type = dnnl::memory::data_type::f32;
|
||||
|
||||
// indicate whether gamma or/and beta are given
|
||||
auto flags =
|
||||
dnnl::normalization_flags::use_global_stats; // don't calculate the mean and variance for each mini-batch
|
||||
if (weights != nullptr) flags |= dnnl::normalization_flags::use_scale_shift;
|
||||
|
||||
dnnl::memory::dims dims;
|
||||
dnnl::memory::format_tag format;
|
||||
|
||||
const sd::LongType indHW = isNCHW ? 2 : 1;
|
||||
const sd::LongType bS = x->sizeAt(0);
|
||||
const sd::LongType iC = isNCHW ? x->sizeAt(1) : x->sizeAt(-1);
|
||||
|
||||
sd::LongType iD, iH, iW;
|
||||
|
||||
if (xRank == 2) {
|
||||
dims = {bS, iC};
|
||||
format = dnnl::memory::format_tag::nc;
|
||||
} else if (xRank == 4) {
|
||||
iH = x->sizeAt(indHW);
|
||||
iW = x->sizeAt(indHW + 1);
|
||||
dims = {bS, iC, iH, iW};
|
||||
format = isNCHW ? dnnl::memory::format_tag::nchw : dnnl::memory::format_tag::nhwc;
|
||||
} else { // xRank = 5
|
||||
iD = x->sizeAt(indHW);
|
||||
iH = x->sizeAt(indHW + 1);
|
||||
iW = x->sizeAt(indHW + 2);
|
||||
dims = {bS, iC, iD, iH, iW};
|
||||
format = isNCHW ? dnnl::memory::format_tag::ncdhw : dnnl::memory::format_tag::ndhwc;
|
||||
}
|
||||
|
||||
// memory descriptors for arrays
|
||||
|
||||
// x
|
||||
dnnl::memory::desc x_mkl_md = dnnl::memory::desc(dims, type, format);
|
||||
dnnl::memory::desc x_user_md = dnnl::memory::desc(dims, type, format);
|
||||
onednnUtils::setBlockStrides(*x, x_user_md);
|
||||
|
||||
// dLdO
|
||||
dnnl::memory::desc dLdO_mkl_md = dnnl::memory::desc(dims, type, dnnl::memory::format_tag::any);
|
||||
dnnl::memory::desc dLdO_user_md = dnnl::memory::desc(dims, type, format);
|
||||
onednnUtils::setBlockStrides(*dLdO, dLdO_user_md);
|
||||
|
||||
// dLdI
|
||||
dnnl::memory::desc dLdI_mkl_md = dnnl::memory::desc(dims, type, dnnl::memory::format_tag::any);
|
||||
dnnl::memory::desc dLdI_user_md = dnnl::memory::desc(dims, type, format);
|
||||
onednnUtils::setBlockStrides(*dLdI, dLdI_user_md);
|
||||
|
||||
auto engine = onednnUtils::getEngine(LaunchContext::defaultContext()->engine());
|
||||
|
||||
// batchnorm forward description
|
||||
dnnl::batch_normalization_forward::desc op_ff_desc(dnnl::prop_kind::forward_inference, x_mkl_md, epsilon, flags);
|
||||
dnnl::batch_normalization_forward::primitive_desc op_ff_prim_desc(op_ff_desc, engine);
|
||||
|
||||
// batchnorm backprop description
|
||||
dnnl::batch_normalization_backward::desc op_bp_desc(dnnl::prop_kind::backward, dLdO_mkl_md, x_mkl_md, epsilon, flags);
|
||||
dnnl::batch_normalization_backward::primitive_desc op_bp_prim_desc(op_bp_desc, engine, op_ff_prim_desc);
|
||||
|
||||
// arguments (memory buffers) necessary for calculations
|
||||
std::unordered_map<int, dnnl::memory> args;
|
||||
|
||||
dnnl::stream stream(engine);
|
||||
|
||||
// provide memory and check whether reorder is required
|
||||
|
||||
// x
|
||||
onednnUtils::loadDataToMklStream(*x, engine, stream, x_user_md, op_bp_prim_desc.src_desc(), args[DNNL_ARG_SRC]);
|
||||
|
||||
// dLdO
|
||||
onednnUtils::loadDataToMklStream(*dLdO, engine, stream, dLdO_user_md, op_bp_prim_desc.diff_dst_desc(),
|
||||
args[DNNL_ARG_DIFF_DST]);
|
||||
|
||||
// mean
|
||||
auto mean_mkl_mem = dnnl::memory(op_bp_prim_desc.mean_desc(), engine, const_cast<void*>(mean->buffer()));
|
||||
args[DNNL_ARG_MEAN] = mean_mkl_mem;
|
||||
|
||||
// variance
|
||||
auto var_mkl_mem = dnnl::memory(op_bp_prim_desc.variance_desc(), engine, const_cast<void*>(variance->buffer()));
|
||||
args[DNNL_ARG_VARIANCE] = var_mkl_mem;
|
||||
|
||||
// dLdI
|
||||
auto dLdI_user_mem = onednnUtils::loadDataToMklStream(*dLdI, engine, stream, dLdI_user_md,
|
||||
op_bp_prim_desc.diff_src_desc(), args[DNNL_ARG_DIFF_SRC]);
|
||||
|
||||
// gamma and beta (and their gradients) if they are present
|
||||
if (weights != nullptr) {
|
||||
auto w_mkl_mem = dnnl::memory(op_bp_prim_desc.weights_desc(), engine, const_cast<void*>(weights->buffer()));
|
||||
args[DNNL_ARG_WEIGHTS] = w_mkl_mem;
|
||||
|
||||
auto dLdW_mkl_mem = dnnl::memory(op_bp_prim_desc.weights_desc(), engine, dLdW->buffer());
|
||||
args[DNNL_ARG_DIFF_WEIGHTS] = dLdW_mkl_mem;
|
||||
}
|
||||
|
||||
// run calculations
|
||||
dnnl::batch_normalization_backward(op_bp_prim_desc).execute(stream, args);
|
||||
|
||||
// reorder outputs if necessary
|
||||
if (op_bp_prim_desc.diff_src_desc() != dLdI_user_mem.get_desc())
|
||||
dnnl::reorder(args[DNNL_ARG_DIFF_SRC], dLdI_user_mem).execute(stream, args[DNNL_ARG_DIFF_SRC], dLdI_user_mem);
|
||||
|
||||
stream.wait();
|
||||
|
||||
|
||||
// notations:
|
||||
// f = g * (gamma * ((x - m) / (v + eps)^0.5) + beta) -> means dLdO * ff_output
|
||||
// g = dLdO
|
||||
// stdInv = 1 / (v + eps)^0.5
|
||||
// N - batch size (product of spatial dimensions)
|
||||
|
||||
// formula for full derivative with respect to input (x)
|
||||
// dLdI = dfdx + dfdm*dmdx + dfdv*(dvdm*dmdx + dvdx)
|
||||
|
||||
// !!! MKL CALCULATES ONLY FIRST TERM dfdx, SO WE SHOULD CALCULATE TERM (dfdm*dmdx + dfdv*(dvdm*dmdx + dvdx)) BY
|
||||
// OURSELF !!!
|
||||
|
||||
// dfdm = -gamma*stdInv*g_sum;
|
||||
// dmdx = 1/N;
|
||||
// dvdx = 2 * (x - m) / N
|
||||
// dvdm = -2 * [(x - m)]_sum / N
|
||||
// dfdv = -0.5 * [g*(x - m)]_sum * stdInv^3, drop gamma here for calc convenience
|
||||
|
||||
// finally:
|
||||
// dLdI = dfdm / N + (2/N) * dfdv * (dvdm/2 + (x - m))
|
||||
// dLdI = gamma * ( stdInv * -g_sum/N + (2/N) * dfdv * (dvdm/2 + (x - m)) )
|
||||
|
||||
std::vector<sd::LongType> axes = isNCHW ? std::vector<sd::LongType>{1} : std::vector<sd::LongType>{xRank - 1};
|
||||
const auto excludedAxes = ShapeUtils::evalDimsToExclude(x->rankOf(),axes.size(), axes.data());
|
||||
|
||||
// inversed batch size 1 / N
|
||||
const auto Ninv = 1.f * mean->lengthOf() / x->lengthOf();
|
||||
|
||||
// x - mean
|
||||
NDArray xMinusMean(x); // empty array with same shape as x
|
||||
const_cast<NDArray*>(x)->applyBroadcast(sd::broadcast::Subtract, &axes, mean, &xMinusMean);
|
||||
|
||||
// stdInv
|
||||
NDArray stdInv = *variance + epsilon;
|
||||
stdInv.applyTransform(transform::Reciprocal, &stdInv); // 1 / (variance + epsilon)
|
||||
stdInv.applyTransform(transform::Sqrt, &stdInv); // 1 / (variance + epsilon)^0.5
|
||||
|
||||
// dfdm / N
|
||||
auto dfdm = dLdO->reduceAlongDimension(sd::reduce::Sum, excludedAxes);
|
||||
dfdm *= stdInv;
|
||||
dfdm *= -Ninv;
|
||||
|
||||
// dvdm / 2
|
||||
NDArray dvdm(mean); // empty array with same shape as mean
|
||||
xMinusMean.reduceAlongDimension(sd::reduce::Sum, &dvdm, excludedAxes);
|
||||
dvdm *= -Ninv;
|
||||
|
||||
// (2/N)*dfdv
|
||||
NDArray dfdv(variance); // empty array with same shape as variance
|
||||
(xMinusMean * *dLdO).reduceAlongDimension(sd::reduce::Sum, &dfdv, excludedAxes);
|
||||
dfdv *= stdInv * stdInv * stdInv;
|
||||
dfdv *= -Ninv;
|
||||
|
||||
// dvdm/2 + (x - m)
|
||||
xMinusMean.applyBroadcast(sd::broadcast::Add, &axes, &dvdm, &xMinusMean);
|
||||
// dfdv * (dvdm/2 + (x - m))
|
||||
xMinusMean.applyBroadcast(sd::broadcast::Multiply, &axes, &dfdv, &xMinusMean);
|
||||
// add dfdm / N
|
||||
xMinusMean.applyBroadcast(sd::broadcast::Add, &axes, &dfdm, &xMinusMean);
|
||||
// * gamma
|
||||
auto gamma = (*weights)({0, 1, 0, 0});
|
||||
xMinusMean.applyBroadcast(sd::broadcast::Multiply, &axes, &gamma, &xMinusMean);
|
||||
|
||||
*dLdI += xMinusMean;
|
||||
}
|
||||
|
||||
PLATFORM_IMPL(batchnorm, ENGINE_CPU) {
|
||||
auto input = INPUT_VARIABLE(0); // 2D:nc, 4D:nchw/nhwc, 5D:ncdhw/ndhwc
|
||||
auto mean = INPUT_VARIABLE(1); // [c]
|
||||
auto variance = INPUT_VARIABLE(2); // [c]
|
||||
NDArray* gamma = nullptr; // [c]
|
||||
NDArray* beta = nullptr; // [c]
|
||||
|
||||
auto output = OUTPUT_VARIABLE(0); // same shape as input
|
||||
|
||||
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 sd::LongType inRank = input->rankOf();
|
||||
|
||||
// get axes args to normalize input array over
|
||||
std::vector<sd::LongType> 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 sd::LongType numOfAxes = axes.size();
|
||||
REQUIRE_TRUE(numOfAxes == 1, 0,
|
||||
"BATCHNORM_MKLDNN op: mkl dnn library supports only one axis which represents channel dimension, but "
|
||||
"got %i axes instead!",
|
||||
numOfAxes);
|
||||
REQUIRE_TRUE(inRank == 2 || inRank == 4 || inRank == 5, 0,
|
||||
"BATCHNORM_MKLDNN op: possible values for rank of input array are 2, 4 or 5, but got %i instead!",
|
||||
inRank);
|
||||
REQUIRE_TRUE(mean->rankOf() == 1 && mean->sizeAt(0) == input->sizeAt(axes[0]), 0,
|
||||
"BATCHNORM_MKLDNN op: wrong shape of mean array, expected is [%lld], but got %s instead !",
|
||||
input->sizeAt(axes[0]), ShapeUtils::shapeAsString(mean).c_str());
|
||||
REQUIRE_TRUE(variance->rankOf() == 1 && variance->sizeAt(0) == input->sizeAt(axes[0]), 0,
|
||||
"BATCHNORM_MKLDNN op: wrong shape of variance array, expected is [%lld], but got %s instead !",
|
||||
input->sizeAt(axes[0]), ShapeUtils::shapeAsString(variance).c_str());
|
||||
if (gamma != nullptr)
|
||||
REQUIRE_TRUE(gamma->rankOf() == 1 && gamma->sizeAt(0) == input->sizeAt(axes[0]), 0,
|
||||
"BATCHNORM_MKLDNN op: wrong shape of gamma array, expected is [%lld], but got %s instead !",
|
||||
input->sizeAt(axes[0]), ShapeUtils::shapeAsString(gamma).c_str());
|
||||
if (beta != nullptr)
|
||||
REQUIRE_TRUE(beta->rankOf() == 1 && beta->sizeAt(0) == input->sizeAt(axes[0]), 0,
|
||||
"BATCHNORM_MKLDNN op: wrong shape of beta array, expected is [%lld], but got %s instead !",
|
||||
input->sizeAt(axes[0]), ShapeUtils::shapeAsString(beta).c_str());
|
||||
|
||||
// types of all input arrays should be the same (except dLdO)
|
||||
for (size_t i = 1; i < block.width() - 1; ++i)
|
||||
REQUIRE_TRUE(INPUT_VARIABLE(0)->dataType() == INPUT_VARIABLE(i)->dataType(), 0,
|
||||
"BATCHNORM_MKLDNN op: types of all input arrays should be the same !");
|
||||
|
||||
NDArray* weights = nullptr;
|
||||
|
||||
if (applyScale || applyOffset) {
|
||||
std::vector<sd::LongType > shape = {2, input->sizeAt(axes[0])};
|
||||
weights = new NDArray(input->ordering(),shape , input->dataType());
|
||||
|
||||
if (applyScale)
|
||||
(*weights)({0, 1, 0, 0}).assign(gamma);
|
||||
else {
|
||||
sd::LongType scalarVal = 1;
|
||||
(*weights)({0, 1, 0, 0}).assign(scalarVal);
|
||||
}
|
||||
if (applyOffset)
|
||||
(*weights)({1, 2, 0, 0}).assign(beta);
|
||||
else
|
||||
(*weights)({1, 2, 0, 0}).assign(0);
|
||||
}
|
||||
|
||||
const bool isNCHW = !(axes[0] == inRank - 1 && inRank > 2);
|
||||
|
||||
batchnormMKLDNN(input, mean, variance, weights, output, epsilon, isNCHW);
|
||||
|
||||
delete weights;
|
||||
|
||||
return sd::Status::OK;
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
PLATFORM_CHECK(batchnorm, ENGINE_CPU) {
|
||||
auto input = INPUT_VARIABLE(0); // 2D:nc, 4D:nchw/nhwc, 5D:ncdhw/ndhwc
|
||||
auto mean = INPUT_VARIABLE(1); // [c]
|
||||
auto variance = INPUT_VARIABLE(2); // [c]
|
||||
NDArray* gamma = nullptr; // [c]
|
||||
NDArray* beta = nullptr; // [c]
|
||||
|
||||
auto output = OUTPUT_VARIABLE(0); // same shape as input
|
||||
|
||||
const bool applyScale = (bool)INT_ARG(0);
|
||||
const bool applyOffset = (bool)INT_ARG(1);
|
||||
|
||||
if (applyScale) gamma = INPUT_VARIABLE(3);
|
||||
if (applyOffset) beta = INPUT_VARIABLE(3 + (int)applyScale);
|
||||
|
||||
const int numOfIntArgs = block.getIArguments()->size();
|
||||
std::vector<sd::LongType> axes;
|
||||
if (numOfIntArgs > 2)
|
||||
for (int i = 2; i < numOfIntArgs; ++i) axes.push_back(INT_ARG(i));
|
||||
else
|
||||
axes.push_back(input->rankOf() - 1); // default dimension to reduce along is last dimension
|
||||
|
||||
const int inRank = input->rankOf();
|
||||
|
||||
Requirements req("ONEDNN BATCHNORM OP");
|
||||
req.expectTrue(block.isUseONEDNN(), IS_USE_ONEDNN_MSG) &&
|
||||
req.expectEq(makeInfoVariable(axes.size(), "axes.size()"), 1) &&
|
||||
req.expectIn(makeInfoVariable(axes[0], "axes#0"), {1, inRank - 1}) &&
|
||||
req.expectIn(makeInfoVariable(inRank, RANK_MSG_INPUT0), {2, 4, 5}) &&
|
||||
req.expectTrue(makeInfoVariable(
|
||||
[input, mean, variance, gamma, beta, output] {
|
||||
DataType inputType = input->dataType();
|
||||
DataType meanType = mean->dataType();
|
||||
DataType varType = variance->dataType();
|
||||
DataType gammaType = gamma != nullptr ? gamma->dataType() : DataType::FLOAT32;
|
||||
DataType betaType = beta != nullptr ? beta->dataType() : DataType::FLOAT32;
|
||||
DataType outType = output->dataType();
|
||||
return (inputType == DataType::FLOAT32 && meanType == DataType::FLOAT32 &&
|
||||
varType == DataType::FLOAT32 && gammaType == DataType::FLOAT32 &&
|
||||
betaType == DataType::FLOAT32 && outType == DataType::FLOAT32);
|
||||
},
|
||||
TYPECHECK_MSG),
|
||||
NO_MSG);
|
||||
req.logTheSuccess();
|
||||
return req;
|
||||
}
|
||||
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
PLATFORM_IMPL(batchnorm_bp, ENGINE_CPU) {
|
||||
NDArray* input = INPUT_VARIABLE(0); // 2D:nc, 4D:nchw/nhwc, 5D:ncdhw/ndhwc
|
||||
NDArray* mean = INPUT_VARIABLE(1); // [c]
|
||||
NDArray* variance = INPUT_VARIABLE(2); // [c]
|
||||
NDArray* gamma = nullptr; // [c]
|
||||
NDArray* beta = nullptr; // [c]
|
||||
NDArray* dLdO = INPUT_VARIABLE(block.width() - 1); // same as input
|
||||
|
||||
NDArray* dLdI = OUTPUT_VARIABLE(0); // same as input
|
||||
NDArray* dLdM = OUTPUT_VARIABLE(1); // [c]
|
||||
NDArray* dLdV = OUTPUT_VARIABLE(2); // [c]
|
||||
NDArray* dLdG = nullptr; // [c]
|
||||
NDArray* dLdB = nullptr; // [c]
|
||||
|
||||
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);
|
||||
dLdG = OUTPUT_VARIABLE(3);
|
||||
}
|
||||
if (applyOffset) {
|
||||
beta = INPUT_VARIABLE(3 + (sd::LongType)applyScale);
|
||||
dLdB = OUTPUT_VARIABLE(3 + (sd::LongType)applyScale);
|
||||
}
|
||||
|
||||
const int numOfIntArgs = block.getIArguments()->size();
|
||||
const int inRank = input->rankOf();
|
||||
|
||||
// get axes args to normalize input array over
|
||||
std::vector<sd::LongType> 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 == 1, 0,
|
||||
"BATCHNORM_BP_MKLDNN op: mkl dnn library supports only one axis which represents channel dimension, but "
|
||||
"got %i axes instead!",
|
||||
numOfAxes);
|
||||
REQUIRE_TRUE(inRank == 2 || inRank == 4 || inRank == 5, 0,
|
||||
"BATCHNORM_BP_MKLDNN op: possible values for rank of input array are 2, 4 or 5, but got %i instead!",
|
||||
inRank);
|
||||
REQUIRE_TRUE(input->isSameShape(dLdO), 0,
|
||||
"BATCHNORM_BP_MKLDNN op: wrong shape of gradients array, expected is %s, but got %s instead !",
|
||||
ShapeUtils::shapeAsString(input).c_str(), ShapeUtils::shapeAsString(dLdO).c_str());
|
||||
REQUIRE_TRUE(mean->rankOf() == 1 && mean->sizeAt(0) == input->sizeAt(axes[0]), 0,
|
||||
"BATCHNORM_BP_MKLDNN op: wrong shape of mean array, expected is [%lld], but got %s instead !",
|
||||
input->sizeAt(axes[0]), ShapeUtils::shapeAsString(mean).c_str());
|
||||
REQUIRE_TRUE(variance->rankOf() == 1 && variance->sizeAt(0) == input->sizeAt(axes[0]), 0,
|
||||
"BATCHNORM_BP_MKLDNN op: wrong shape of variance array, expected is [%lld], but got %s instead !",
|
||||
input->sizeAt(axes[0]), ShapeUtils::shapeAsString(variance).c_str());
|
||||
if (gamma != nullptr)
|
||||
REQUIRE_TRUE(gamma->rankOf() == 1 && gamma->sizeAt(0) == input->sizeAt(axes[0]), 0,
|
||||
"BATCHNORM_BP_MKLDNN op: wrong shape of gamma array, expected is [%lld], but got %s instead !",
|
||||
input->sizeAt(axes[0]), ShapeUtils::shapeAsString(gamma).c_str());
|
||||
if (beta != nullptr)
|
||||
REQUIRE_TRUE(beta->rankOf() == 1 && beta->sizeAt(0) == input->sizeAt(axes[0]), 0,
|
||||
"BATCHNORM_BP_MKLDNN op: wrong shape of beta array, expected is [%lld], but got %s instead !",
|
||||
input->sizeAt(axes[0]), ShapeUtils::shapeAsString(beta).c_str());
|
||||
|
||||
// types of all input arrays should be the same
|
||||
for (size_t i = 1; i < block.width() - 1; ++i)
|
||||
REQUIRE_TRUE(INPUT_VARIABLE(0)->dataType() == INPUT_VARIABLE(i)->dataType(), 0,
|
||||
"BATCHNORM_BP_MKLDNN op: types of all input arrays should be the same !");
|
||||
|
||||
NDArray *weights = nullptr, *dLdW = nullptr;
|
||||
|
||||
if (applyScale || applyOffset) {
|
||||
sd::LongType scalar = 1;
|
||||
std::vector<sd::LongType> shape = {2, input->sizeAt(axes[0])};
|
||||
weights = new NDArray(input->ordering(),shape, input->dataType());
|
||||
dLdW = new NDArray(input->ordering(), shape, input->dataType());
|
||||
if (applyScale)
|
||||
(*weights)({0, 1, 0, 0}).assign(gamma);
|
||||
else
|
||||
(*weights)({0, 1, 0, 0}).assign(scalar);
|
||||
if (applyOffset)
|
||||
(*weights)({1, 2, 0, 0}).assign(beta);
|
||||
else
|
||||
(*weights)({1, 2, 0, 0}).assign(0);
|
||||
}
|
||||
|
||||
const bool isNCHW = !(axes[0] == inRank - 1 && inRank > 2);
|
||||
|
||||
if (shape::strideDescendingCAscendingF(dLdO->shapeInfo()))
|
||||
batchnormBpMKLDNN(input, mean, variance, dLdO, weights, dLdI, dLdW, epsilon, isNCHW);
|
||||
else {
|
||||
NDArray dupped = dLdO->dup();
|
||||
batchnormBpMKLDNN(input, mean, variance, &dupped, weights, dLdI, dLdW, epsilon, isNCHW);
|
||||
}
|
||||
*dLdM = 0;
|
||||
*dLdV = 0;
|
||||
|
||||
if (applyScale || applyOffset) {
|
||||
if (applyScale) {
|
||||
NDArray assign = (*dLdW)({0, 1, 0, 0});
|
||||
dLdG->assign(&assign);
|
||||
}
|
||||
if (applyOffset) {
|
||||
NDArray assign = (*dLdW)({1, 2, 0, 0});
|
||||
dLdB->assign(&assign);
|
||||
}
|
||||
|
||||
delete weights;
|
||||
delete dLdW;
|
||||
}
|
||||
|
||||
return sd::Status::OK;
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
PLATFORM_CHECK(batchnorm_bp, ENGINE_CPU) {
|
||||
NDArray* input = INPUT_VARIABLE(0); // 2D:nc, 4D:nchw, 5D:ncdhw
|
||||
NDArray* mean = INPUT_VARIABLE(1); // [c]
|
||||
NDArray* variance = INPUT_VARIABLE(2); // [c]
|
||||
NDArray* dLdO = INPUT_VARIABLE(3); // same as input
|
||||
NDArray* gamma = nullptr; // [c]
|
||||
NDArray* beta = nullptr; // [c]
|
||||
|
||||
NDArray* dLdI = OUTPUT_VARIABLE(0); // same as input
|
||||
NDArray* dLdM = OUTPUT_VARIABLE(1); // [c]
|
||||
NDArray* dLdV = OUTPUT_VARIABLE(2); // [c]
|
||||
NDArray* dLdG = nullptr; // [c]
|
||||
NDArray* dLdB = nullptr; // [c]
|
||||
|
||||
const bool applyScale = (bool)INT_ARG(0);
|
||||
const bool applyOffset = (bool)INT_ARG(1);
|
||||
|
||||
if (applyScale) {
|
||||
gamma = INPUT_VARIABLE(4);
|
||||
dLdG = OUTPUT_VARIABLE(3);
|
||||
}
|
||||
if (applyOffset) {
|
||||
beta = INPUT_VARIABLE(4 + (sd::LongType)applyScale);
|
||||
dLdB = OUTPUT_VARIABLE(3 + (sd::LongType)applyScale);
|
||||
}
|
||||
|
||||
const int numOfIntArgs = block.getIArguments()->size();
|
||||
std::vector<sd::LongType> axes;
|
||||
if (numOfIntArgs > 2)
|
||||
for (int i = 2; i < numOfIntArgs; ++i) axes.push_back(INT_ARG(i));
|
||||
else
|
||||
axes.push_back(input->rankOf() - 1); // default dimension to reduce along is last dimension
|
||||
|
||||
const sd::LongType inRank = input->rankOf();
|
||||
std::vector<sd::LongType> shape = {1, inRank - 1};
|
||||
Requirements req("ONEDNN BATCHNORM_BP OP");
|
||||
req.expectTrue(block.isUseONEDNN(), IS_USE_ONEDNN_MSG) &&
|
||||
req.expectEq(makeInfoVariable(axes.size(), "axes.size()"), 1) &&
|
||||
req.expectIn(makeInfoVariable(inRank, RANK_MSG_INPUT0), {2, 4, 5}) &&
|
||||
req.expectTrue(makeInfoVariable(
|
||||
[input, mean, variance, dLdO, gamma, beta, dLdG, dLdB, dLdI] {
|
||||
DataType inputType = input->dataType();
|
||||
DataType meanType = mean->dataType();
|
||||
DataType varType = variance->dataType();
|
||||
DataType dLdOType = dLdO->dataType();
|
||||
DataType gammaType = gamma != nullptr ? gamma->dataType() : DataType::FLOAT32;
|
||||
DataType betaType = beta != nullptr ? beta->dataType() : DataType::FLOAT32;
|
||||
|
||||
DataType dLdIType = dLdI->dataType();
|
||||
DataType dLdGType = gamma != nullptr ? dLdG->dataType() : DataType::FLOAT32;
|
||||
DataType dLdBType = beta != nullptr ? dLdB->dataType() : DataType::FLOAT32;
|
||||
return (inputType == DataType::FLOAT32 && meanType == DataType::FLOAT32 &&
|
||||
varType == DataType::FLOAT32 && dLdOType == DataType::FLOAT32 &&
|
||||
gammaType == DataType::FLOAT32 && betaType == DataType::FLOAT32 &&
|
||||
dLdIType == DataType::FLOAT32 && dLdGType == DataType::FLOAT32 &&
|
||||
dLdBType == DataType::FLOAT32);
|
||||
},
|
||||
TYPECHECK_MSG),
|
||||
NO_MSG);
|
||||
req.logTheSuccess();
|
||||
return req;
|
||||
}
|
||||
|
||||
} // namespace platforms
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,197 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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 <helpers/MKLDNNStream.h>
|
||||
#include <ops/declarable/OpRegistrator.h>
|
||||
#include <ops/declarable/PlatformHelper.h>
|
||||
#include <system/platform_boilerplate.h>
|
||||
|
||||
#include <numeric>
|
||||
|
||||
#include "mkldnnUtils.h"
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace platforms {
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
static void concatMKLDNN(const std::vector<NDArray*>& inArrs, NDArray& output, const int axis) {
|
||||
// data type
|
||||
dnnl::memory::data_type type;
|
||||
if (output.dataType() == DataType::FLOAT32)
|
||||
type = dnnl::memory::data_type::f32;
|
||||
else if (output.dataType() == DataType::HALF)
|
||||
type = dnnl::memory::data_type::f16;
|
||||
else if (output.dataType() == DataType::BFLOAT16)
|
||||
type = dnnl::memory::data_type::bf16;
|
||||
else if (output.dataType() == DataType::UINT8)
|
||||
type = dnnl::memory::data_type::u8;
|
||||
else
|
||||
type = dnnl::memory::data_type::s8;
|
||||
|
||||
std::vector<dnnl::memory::desc> x_user_md(inArrs.size()), x_mkl_md(inArrs.size());
|
||||
|
||||
// inputs
|
||||
for (size_t i = 0; i < inArrs.size(); ++i) {
|
||||
dnnl::memory::dims dims = inArrs[i]->getShapeAsFlatVector();
|
||||
x_user_md[i] = x_mkl_md[i] = dnnl::memory::desc(dims, type, onednnUtils::getFormat(*inArrs[i]));
|
||||
onednnUtils::setBlockStrides(*inArrs[i], x_user_md[i]);
|
||||
}
|
||||
|
||||
// output
|
||||
dnnl::memory::dims dims = output.getShapeAsFlatVector();
|
||||
dnnl::memory::desc z_mkl_md = dnnl::memory::desc(dims, type, dnnl::memory::format_tag::any);
|
||||
dnnl::memory::desc z_user_md = dnnl::memory::desc(dims, type, onednnUtils::getFormat(output));
|
||||
onednnUtils::setBlockStrides(output, z_user_md);
|
||||
|
||||
std::unordered_map<int, dnnl::memory> args;
|
||||
|
||||
auto engine = onednnUtils::getEngine(LaunchContext::defaultContext()->engine());
|
||||
|
||||
dnnl::concat::primitive_desc op_prim_desc(axis, x_mkl_md, engine);
|
||||
|
||||
dnnl::stream stream(engine);
|
||||
|
||||
// inputs
|
||||
for (size_t i = 0; i < inArrs.size(); ++i)
|
||||
onednnUtils::loadDataToMklStream(*inArrs[i], engine, stream, x_user_md[i], op_prim_desc.src_desc(i),
|
||||
args[DNNL_ARG_MULTIPLE_SRC + i]);
|
||||
|
||||
// outputs
|
||||
auto z_user_mem =
|
||||
onednnUtils::loadDataToMklStream(output, engine, stream, z_user_md, op_prim_desc.dst_desc(), args[DNNL_ARG_DST]);
|
||||
|
||||
// primitive execution
|
||||
dnnl::concat(op_prim_desc).execute(stream, args);
|
||||
|
||||
// reorder output if necessary
|
||||
if (op_prim_desc.dst_desc() != z_user_mem.get_desc())
|
||||
dnnl::reorder(args[DNNL_ARG_DST], z_user_mem).execute(stream, args[DNNL_ARG_DST], z_user_mem);
|
||||
|
||||
stream.wait();
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
PLATFORM_IMPL(concat, ENGINE_CPU) {
|
||||
REQUIRE_TRUE(block.width() > 0, 0, "CONCAT MKLDNN op: No input arrays were provided");
|
||||
|
||||
const bool isAxisInLastArr = block.getBArguments()->size() == 0 ? false : B_ARG(0);
|
||||
|
||||
const int numOfInArrs = isAxisInLastArr ? block.width() - 1 : block.width();
|
||||
|
||||
// first of all take into account possible presence of empty arrays
|
||||
// also if scalar is present -> copy its value to vector with length=1
|
||||
std::vector<NDArray*> nonEmptyArrs;
|
||||
std::vector<sd::LongType> arrsToDelete;
|
||||
int index = 0;
|
||||
bool allOfSameType = true;
|
||||
auto rankOfFirstArr = block.width() > 0 ? INPUT_VARIABLE(0)->rankOf() : 0;
|
||||
auto typeOfFirstArr = block.width() > 0 ? INPUT_VARIABLE(0)->dataType() : block.dataType();
|
||||
|
||||
for (int i = 0; i < numOfInArrs; ++i) {
|
||||
auto input = INPUT_VARIABLE(i);
|
||||
auto currentRank = input->rankOf();
|
||||
|
||||
if (!input->isEmpty()) {
|
||||
allOfSameType &= (typeOfFirstArr == input->dataType());
|
||||
|
||||
if (input->rankOf() == 0) {
|
||||
std::vector<sd::LongType> dim = {1};
|
||||
auto vec = new NDArray('c', dim, input->dataType(), block.launchContext());
|
||||
vec->assign(input);
|
||||
nonEmptyArrs.push_back(vec);
|
||||
arrsToDelete.push_back(index);
|
||||
} else {
|
||||
nonEmptyArrs.push_back(input);
|
||||
}
|
||||
++index;
|
||||
}
|
||||
}
|
||||
|
||||
const int numOfNonEmptyArrs = nonEmptyArrs.size();
|
||||
|
||||
if (numOfNonEmptyArrs == 0) {
|
||||
// All inputs are empty arrays -> return empty, mainly for TF import compatibility (no op)
|
||||
REQUIRE_TRUE(OUTPUT_VARIABLE(0)->isEmpty(), 0,
|
||||
"CONCAT MKLDNN op: If all input variables are empty, output must be empty");
|
||||
return sd::Status::OK;
|
||||
}
|
||||
|
||||
const int rank = nonEmptyArrs[0]->rankOf(); // look up to first non-empty array
|
||||
int axis = isAxisInLastArr ? INPUT_VARIABLE(block.width() - 1)->e<int>(0) : INT_ARG(0);
|
||||
if (axis < 0) {
|
||||
axis += rank;
|
||||
}
|
||||
|
||||
// ******** input validation ******** //
|
||||
REQUIRE_TRUE(allOfSameType, 0, "CONCAT MKLDNN op: all of input arrays must have same type !");
|
||||
REQUIRE_TRUE(nonEmptyArrs[0]->dataType() == OUTPUT_VARIABLE(0)->dataType(), 0,
|
||||
"CONCAT MKLDNN op: output array should have the same type as inputs arrays !");
|
||||
REQUIRE_TRUE(0 <= axis && (axis < rank || (axis == 0 && rank == 0)), 0,
|
||||
"CONCAT MKLDNN op: input axis must be in range [0, %i], but got %i instead!", rank - 1, axis);
|
||||
|
||||
for (int i = 1; i < numOfNonEmptyArrs; ++i)
|
||||
REQUIRE_TRUE(nonEmptyArrs[i]->rankOf() == rank, 0, "CONCAT MKLDNN op: all input arrays must have the same rank !");
|
||||
|
||||
for (int i = 1; i < numOfNonEmptyArrs; ++i) {
|
||||
for (int dim = 0; dim < rank; ++dim)
|
||||
if (dim != axis)
|
||||
REQUIRE_TRUE(nonEmptyArrs[i]->sizeAt(dim) == nonEmptyArrs[0]->sizeAt(dim), 0,
|
||||
"CONCAT MKLDNN op: all input arrays must have the same dimensions (except those on input axis) !");
|
||||
}
|
||||
// ******** end of input validation ******** //
|
||||
|
||||
auto output = OUTPUT_VARIABLE(0);
|
||||
|
||||
if (numOfNonEmptyArrs == 1)
|
||||
output->assign(nonEmptyArrs[0]);
|
||||
else
|
||||
concatMKLDNN(nonEmptyArrs, *output, axis);
|
||||
|
||||
|
||||
return sd::Status::OK;
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
PLATFORM_CHECK(concat, ENGINE_CPU) {
|
||||
auto z = OUTPUT_VARIABLE(0);
|
||||
|
||||
const bool isAxisInLastArr = block.getBArguments()->size() == 0 ? false : B_ARG(0);
|
||||
const int numOfInArrs = isAxisInLastArr ? block.width() - 1 : block.width();
|
||||
Requirements req("ONEDNN CONCAT OP");
|
||||
req.expectTrue(block.isUseONEDNN(), IS_USE_ONEDNN_MSG) &&
|
||||
req.expectLess(makeInfoVariable(z->rankOf(), RANK_MSG_OUTPUT), 7) &&
|
||||
req.expectLessEq(makeInfoVariable(numOfInArrs, "numOfinArrs"), 3072) &&
|
||||
req.expectTrue(makeInfoVariable(
|
||||
[z] {
|
||||
const auto zType = z->dataType();
|
||||
return (zType == DataType::FLOAT32 || zType == DataType::HALF ||
|
||||
zType == DataType::BFLOAT16 || zType == DataType::UINT8 || zType == DataType::INT8);
|
||||
},
|
||||
TYPECHECK_MSG),
|
||||
NO_MSG);
|
||||
req.logTheSuccess();
|
||||
return req;
|
||||
}
|
||||
|
||||
} // namespace platforms
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,439 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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 saudet
|
||||
// @author raver119@gmail.com
|
||||
// @author Yurii Shyrma (iuriish@yahoo.com)
|
||||
//
|
||||
#include <helpers/MKLDNNStream.h>
|
||||
#include <ops/declarable/OpRegistrator.h>
|
||||
#include <ops/declarable/PlatformHelper.h>
|
||||
#include <ops/declarable/helpers/convolutions.h>
|
||||
#include <system/platform_boilerplate.h>
|
||||
|
||||
#include "mkldnnUtils.h"
|
||||
|
||||
using namespace dnnl;
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace platforms {
|
||||
|
||||
//////////////////////////////////////////////////////////////////////
|
||||
static void conv2dMKLDNN(NDArray *input, NDArray *weights, NDArray *bias, NDArray *output,
|
||||
const sd::LongType kH, const sd::LongType kW, const sd::LongType sH, const sd::LongType sW, const sd::LongType pH, const sd::LongType pW,
|
||||
const sd::LongType dH, const sd::LongType dW, const int paddingMode, const int isNCHW, const int wFormat) {
|
||||
// mkl support weights in [oC, iC, kH, kW] format only
|
||||
|
||||
sd::LongType bS, iC, iH, iW, oC, oH,
|
||||
oW; // batch size, input channels, input height/width, output channels, output height/width;
|
||||
sd::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);
|
||||
|
||||
sd_debug("Running conv2d onednn with strides: %d %d padding: %d %d dilation: %d %d paddingMode %d weightFormat %d\n",sH,sW,pH,pW,dH,dW,paddingMode,wFormat);
|
||||
const int pWSame = (paddingMode == 2 && dW > 1) ? ((oW - 1) * sW + (kW - 1) * dW + 1 - iW) / 2
|
||||
: pW; // dH == 1 for causal mode in conv1d
|
||||
|
||||
dnnl::memory::dims strides = {sH, sW};
|
||||
dnnl::memory::dims padding = {pH, pW};
|
||||
dnnl::memory::dims padding_r = {(oH - 1) * sH - iH + kH - pH, (oW - 1) * sW - iW + kW - pWSame};
|
||||
dnnl::memory::dims dilation = {dH - 1, dW - 1};
|
||||
|
||||
auto xzFormatMkl = isNCHW ? dnnl::memory::format_tag::nchw : dnnl::memory::format_tag::nhwc;
|
||||
dnnl::memory::format_tag wFormatMkl = dnnl::memory::format_tag::oihw;
|
||||
|
||||
dnnl::memory::dims xDims = {bS, iC, iH, iW};
|
||||
dnnl::memory::dims wDims = {oC, iC, kH, kW};
|
||||
dnnl::memory::dims zDims = {bS, oC, oH, oW};
|
||||
|
||||
auto type = dnnl::memory::data_type::f32;
|
||||
|
||||
std::vector<int> permut;
|
||||
if (0 == wFormat)
|
||||
permut = {3, 2, 0, 1}; // [kH, kW, iC, oC] -> [oC, iC, kH, kW]
|
||||
else if (2 == wFormat)
|
||||
permut = {0, 3, 1, 2}; // [oC, kH, kW, iC] -> [oC, iC, kH, kW]
|
||||
|
||||
// memory descriptors for arrays
|
||||
|
||||
sd_debug("Creating input descriptor\n",0);
|
||||
// input
|
||||
dnnl::memory::desc x_mkl_md = dnnl::memory::desc(xDims, type, dnnl::memory::format_tag::any);
|
||||
dnnl::memory::desc x_user_md = dnnl::memory::desc(xDims, type, xzFormatMkl);
|
||||
onednnUtils::setBlockStrides(*input, x_user_md);
|
||||
|
||||
sd_debug("Creating weight descriptor\n",0);
|
||||
|
||||
// weights
|
||||
dnnl::memory::desc w_mkl_md = dnnl::memory::desc(wDims, type, dnnl::memory::format_tag::any);
|
||||
dnnl::memory::desc w_user_md = dnnl::memory::desc(wDims, type, wFormatMkl);
|
||||
onednnUtils::setBlockStrides(*weights, w_user_md, permut);
|
||||
|
||||
sd_debug("Creating bias descriptor\n",0);
|
||||
|
||||
// bias
|
||||
dnnl::memory::desc b_mkl_md;
|
||||
if (bias != nullptr) b_mkl_md = dnnl::memory::desc({oC}, type, dnnl::memory::format_tag::x);
|
||||
|
||||
sd_debug("Creating output\n",0);
|
||||
|
||||
// output
|
||||
dnnl::memory::desc z_mkl_md = dnnl::memory::desc(zDims, type, dnnl::memory::format_tag::any);
|
||||
dnnl::memory::desc z_user_md = dnnl::memory::desc(zDims, type, xzFormatMkl);
|
||||
onednnUtils::setBlockStrides(*output, z_user_md);
|
||||
|
||||
auto engine = onednnUtils::getEngine(LaunchContext::defaultContext()->engine());
|
||||
|
||||
sd_debug("Creating op descriptor\n",0);
|
||||
|
||||
// operation primitive description
|
||||
dnnl::convolution_forward::desc op_desc(dnnl::prop_kind::forward_inference, dnnl::algorithm::convolution_auto,
|
||||
x_mkl_md, w_mkl_md, b_mkl_md, z_mkl_md, strides, dilation, padding,
|
||||
padding_r);
|
||||
|
||||
sd_debug("Creating prim descriptor\n",0);
|
||||
|
||||
dnnl::convolution_forward::primitive_desc op_prim_desc(op_desc, engine);
|
||||
sd_debug("Created engine\n",0);
|
||||
|
||||
// arguments (memory buffers) necessary for calculations
|
||||
std::unordered_map<int, dnnl::memory> args;
|
||||
|
||||
dnnl::stream stream(engine);
|
||||
|
||||
// provide memory buffers and check whether reorder is required
|
||||
|
||||
// input
|
||||
onednnUtils::loadDataToMklStream(*input, engine, stream, x_user_md, op_prim_desc.src_desc(), args[DNNL_ARG_SRC]);
|
||||
|
||||
// weights
|
||||
onednnUtils::loadDataToMklStream(*weights, engine, stream, w_user_md, op_prim_desc.weights_desc(),
|
||||
args[DNNL_ARG_WEIGHTS]);
|
||||
|
||||
// bias
|
||||
if (bias != nullptr) {
|
||||
auto b_mkl_mem = dnnl::memory(b_mkl_md, engine, const_cast<void *>(bias->buffer()));
|
||||
args[DNNL_ARG_BIAS] = b_mkl_mem;
|
||||
}
|
||||
|
||||
// output
|
||||
auto z_user_mem =
|
||||
onednnUtils::loadDataToMklStream(*output, engine, stream, z_user_md, op_prim_desc.dst_desc(), args[DNNL_ARG_DST]);
|
||||
|
||||
// run calculations
|
||||
dnnl::convolution_forward(op_prim_desc).execute(stream, args);
|
||||
|
||||
// reorder outputs if necessary
|
||||
if (op_prim_desc.dst_desc() != z_user_mem.get_desc())
|
||||
dnnl::reorder(args[DNNL_ARG_DST], z_user_mem).execute(stream, args[DNNL_ARG_DST], z_user_mem);
|
||||
|
||||
stream.wait();
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////
|
||||
static void conv2dBpMKLDNN(NDArray *input, NDArray *weights, NDArray *bias, NDArray *gradO,
|
||||
NDArray *gradI, NDArray *gradW, NDArray *gradB, 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 int paddingMode,
|
||||
const int isNCHW, const int wFormat) {
|
||||
// mkl support weights/gradW in [oC, iC, kH, kW] format only
|
||||
|
||||
sd::LongType bS, iC, iH, iW, oC, oH,
|
||||
oW; // batch size, input channels, input height/width, output channels, output height/width;
|
||||
sd::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);
|
||||
|
||||
const int pWSame = (paddingMode == 2 && dW > 1) ? ((oW - 1) * sW + (kW - 1) * dW + 1 - iW) / 2
|
||||
: pW; // dH == 1 for causal mode in conv1d
|
||||
|
||||
dnnl::memory::dims strides = {sH, sW};
|
||||
dnnl::memory::dims padding = {pH, pW};
|
||||
dnnl::memory::dims padding_r = {(oH - 1) * sH - iH + kH - pH, (oW - 1) * sW - iW + kW - pWSame};
|
||||
dnnl::memory::dims dilation = {dH - 1, dW - 1};
|
||||
|
||||
auto xzFormatMkl = isNCHW ? dnnl::memory::format_tag::nchw : dnnl::memory::format_tag::nhwc;
|
||||
dnnl::memory::format_tag wFormatMkl = dnnl::memory::format_tag::oihw;
|
||||
|
||||
dnnl::memory::dims xDims = {bS, iC, iH, iW};
|
||||
dnnl::memory::dims wDims = {oC, iC, kH, kW};
|
||||
dnnl::memory::dims zDims = {bS, oC, oH, oW};
|
||||
|
||||
auto type = dnnl::memory::data_type::f32;
|
||||
|
||||
std::vector<int> permut;
|
||||
if (0 == wFormat)
|
||||
permut = {3, 2, 0, 1}; // [kH, kW, iC, oC] -> [oC, iC, kH, kW]
|
||||
else if (2 == wFormat)
|
||||
permut = {0, 3, 1, 2}; // [oC, kH, kW, iC] -> [oC, iC, kH, kW]
|
||||
|
||||
// memory descriptors for arrays
|
||||
|
||||
// input
|
||||
dnnl::memory::desc x_mkl_md = dnnl::memory::desc(xDims, type, dnnl::memory::format_tag::any);
|
||||
dnnl::memory::desc x_user_md = dnnl::memory::desc(xDims, type, xzFormatMkl);
|
||||
onednnUtils::setBlockStrides(*input, x_user_md);
|
||||
|
||||
// weights
|
||||
dnnl::memory::desc w_mkl_md = dnnl::memory::desc(wDims, type, dnnl::memory::format_tag::any);
|
||||
dnnl::memory::desc w_user_md = dnnl::memory::desc(wDims, type, wFormatMkl);
|
||||
onednnUtils::setBlockStrides(*weights, w_user_md, permut);
|
||||
|
||||
// gradO
|
||||
dnnl::memory::desc gradO_mkl_md = dnnl::memory::desc(zDims, type, dnnl::memory::format_tag::any);
|
||||
dnnl::memory::desc gradO_user_md = dnnl::memory::desc(zDims, type, xzFormatMkl);
|
||||
onednnUtils::setBlockStrides(*gradO, gradO_user_md);
|
||||
|
||||
// gradI
|
||||
dnnl::memory::desc gradI_mkl_md = dnnl::memory::desc(xDims, type, dnnl::memory::format_tag::any);
|
||||
dnnl::memory::desc gradI_user_md = dnnl::memory::desc(xDims, type, xzFormatMkl);
|
||||
onednnUtils::setBlockStrides(*gradI, gradI_user_md);
|
||||
|
||||
// gradW
|
||||
dnnl::memory::desc gradW_mkl_md = dnnl::memory::desc(wDims, type, dnnl::memory::format_tag::any);
|
||||
dnnl::memory::desc gradW_user_md = dnnl::memory::desc(wDims, type, wFormatMkl);
|
||||
onednnUtils::setBlockStrides(*gradW, gradW_user_md, permut);
|
||||
|
||||
// gradB
|
||||
dnnl::memory::desc gradB_mkl_md;
|
||||
if (gradB != nullptr) gradB_mkl_md = dnnl::memory::desc({oC}, type, dnnl::memory::format_tag::x);
|
||||
|
||||
auto engine = onednnUtils::getEngine(LaunchContext::defaultContext()->engine());
|
||||
|
||||
// forward primitive description
|
||||
dnnl::convolution_forward::desc op_ff_desc(dnnl::prop_kind::forward_inference, dnnl::algorithm::convolution_auto,
|
||||
x_mkl_md, w_mkl_md, gradB_mkl_md, gradO_mkl_md, strides, dilation, padding,
|
||||
padding_r);
|
||||
dnnl::convolution_forward::primitive_desc op_ff_prim_desc(op_ff_desc, engine);
|
||||
|
||||
// backward data primitive description
|
||||
dnnl::convolution_backward_data::desc op_data_bp_desc(dnnl::algorithm::convolution_auto, gradI_mkl_md, w_mkl_md,
|
||||
gradO_mkl_md, strides, dilation, padding, padding_r);
|
||||
dnnl::convolution_backward_data::primitive_desc op_data_bp_prim_desc(op_data_bp_desc, engine, op_ff_prim_desc);
|
||||
|
||||
// backward weights primitive description
|
||||
dnnl::convolution_backward_weights::desc op_weights_bp_desc(dnnl::algorithm::convolution_auto, x_mkl_md, gradW_mkl_md,
|
||||
gradB_mkl_md, gradO_mkl_md, strides, dilation, padding,
|
||||
padding_r);
|
||||
dnnl::convolution_backward_weights::primitive_desc op_weights_bp_prim_desc(op_weights_bp_desc, engine,
|
||||
op_ff_prim_desc);
|
||||
|
||||
// arguments (memory buffers) necessary for calculations
|
||||
std::unordered_map<int, dnnl::memory> args;
|
||||
|
||||
dnnl::stream stream(engine);
|
||||
|
||||
// provide memory buffers and check whether reorder is required
|
||||
|
||||
// input
|
||||
onednnUtils::loadDataToMklStream(*input, engine, stream, x_user_md, op_weights_bp_prim_desc.src_desc(),
|
||||
args[DNNL_ARG_SRC]);
|
||||
|
||||
// weights
|
||||
onednnUtils::loadDataToMklStream(*weights, engine, stream, w_user_md, op_data_bp_prim_desc.weights_desc(),
|
||||
args[DNNL_ARG_WEIGHTS]);
|
||||
|
||||
// gradO
|
||||
auto gradO_user_mem = dnnl::memory(gradO_user_md, engine, const_cast<void *>(gradO->buffer()));
|
||||
const bool gradOReorderW = op_weights_bp_prim_desc.diff_dst_desc() != gradO_user_mem.get_desc();
|
||||
const bool gradOReorderD = op_data_bp_prim_desc.diff_dst_desc() != gradO_user_mem.get_desc();
|
||||
auto gradO_mkl_memW = gradOReorderW ? dnnl::memory(op_weights_bp_prim_desc.diff_dst_desc(), engine) : gradO_user_mem;
|
||||
auto gradO_mkl_memD = gradOReorderD ? dnnl::memory(op_data_bp_prim_desc.diff_dst_desc(), engine) : gradO_user_mem;
|
||||
if (gradOReorderW) dnnl::reorder(gradO_user_mem, gradO_mkl_memW).execute(stream, gradO_user_mem, gradO_mkl_memW);
|
||||
if (gradOReorderD) dnnl::reorder(gradO_user_mem, gradO_mkl_memD).execute(stream, gradO_user_mem, gradO_mkl_memD);
|
||||
args[DNNL_ARG_DIFF_DST] = gradO_mkl_memD;
|
||||
|
||||
// gradI
|
||||
auto gradI_user_mem = onednnUtils::loadDataToMklStream(*gradI, engine, stream, gradI_user_md,
|
||||
op_data_bp_prim_desc.diff_src_desc(), args[DNNL_ARG_DIFF_SRC]);
|
||||
|
||||
// gradW
|
||||
auto gradW_user_mem = onednnUtils::loadDataToMklStream(
|
||||
*gradW, engine, stream, gradW_user_md, op_weights_bp_prim_desc.diff_weights_desc(), args[DNNL_ARG_DIFF_WEIGHTS]);
|
||||
|
||||
// gradB
|
||||
if (gradB != nullptr) {
|
||||
auto gradB_mkl_mem = dnnl::memory(gradB_mkl_md, engine, gradB->buffer());
|
||||
args[DNNL_ARG_DIFF_BIAS] = gradB_mkl_mem;
|
||||
}
|
||||
|
||||
// run backward data calculations
|
||||
dnnl::convolution_backward_data(op_data_bp_prim_desc).execute(stream, args);
|
||||
|
||||
if (gradOReorderW || gradOReorderD) args[DNNL_ARG_DIFF_DST] = gradO_mkl_memW;
|
||||
|
||||
// run backward weights calculations
|
||||
dnnl::convolution_backward_weights(op_weights_bp_prim_desc).execute(stream, args);
|
||||
|
||||
// reorder gradI if necessary
|
||||
if (op_data_bp_prim_desc.diff_src_desc() != gradI_user_mem.get_desc())
|
||||
dnnl::reorder(args[DNNL_ARG_DIFF_SRC], gradI_user_mem).execute(stream, args[DNNL_ARG_DIFF_SRC], gradI_user_mem);
|
||||
if (op_weights_bp_prim_desc.diff_weights_desc() != gradW_user_mem.get_desc())
|
||||
dnnl::reorder(args[DNNL_ARG_DIFF_WEIGHTS], gradW_user_mem)
|
||||
.execute(stream, args[DNNL_ARG_DIFF_WEIGHTS], gradW_user_mem);
|
||||
|
||||
stream.wait();
|
||||
|
||||
}
|
||||
|
||||
|
||||
|
||||
//////////////////////////////////////////////////////////////////////
|
||||
PLATFORM_IMPL(conv2d, ENGINE_CPU) {
|
||||
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)
|
||||
|
||||
sd::LongType sH = INT_ARG(2); // strides height
|
||||
sd::LongType sW = INT_ARG(3); // strides width
|
||||
sd::LongType pH = INT_ARG(4); // paddings height
|
||||
sd::LongType pW = INT_ARG(5); // paddings width
|
||||
sd::LongType dH = INT_ARG(6); // dilations height
|
||||
sd::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]
|
||||
|
||||
sd::LongType kH = INT_ARG(0) > 0 ? INT_ARG(0) : static_cast<sd::LongType>(weights->sizeAt(0)); // filter(kernel) height
|
||||
sd::LongType kW = INT_ARG(1) > 0 ? INT_ARG(1) : static_cast<sd::LongType>(weights->sizeAt(1)); // filter(kernel) width
|
||||
|
||||
sd::LongType bS, iC, iH, iW, oC, oH,
|
||||
oW; // batch size, input channels, input height/width, output channels, output height/width;
|
||||
sd::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<sd::LongType> expectedWeightsShape = ConvolutionUtils::expectWeightsShape(wFormat, kH, kW, iC, oC);
|
||||
REQUIRE_TRUE(weights->isSameShape(expectedWeightsShape), 0,
|
||||
"CONV2D MKLDNN 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,
|
||||
"CONV2D MKLDNN OP: wrong shape of array with biases, expected rank, length: <=2, %i, but got %i, %i instead !",
|
||||
oC, bias->rankOf(), bias->lengthOf());
|
||||
|
||||
conv2dMKLDNN(input, weights, bias, output, kH, kW, sH, sW, pH, pW, dH, dW, paddingMode, isNCHW, wFormat);
|
||||
|
||||
return sd::Status::OK;
|
||||
}
|
||||
|
||||
PLATFORM_CHECK(conv2d, ENGINE_CPU) {
|
||||
auto input = INPUT_VARIABLE(0);
|
||||
auto weights = INPUT_VARIABLE(1);
|
||||
|
||||
// conv2d is only available for float32 dtype
|
||||
Requirements req("ONEDNN CONV2d OP");
|
||||
req.expectTrue(block.isUseONEDNN(), IS_USE_ONEDNN_MSG) &&
|
||||
req.expectEq(makeInfoVariable(input->dataType(), TYPE_MSG_INPUT0), sd::DataType::FLOAT32) &&
|
||||
req.expectEq(makeInfoVariable(weights->dataType(), TYPE_MSG_INPUT1), sd::DataType::FLOAT32);
|
||||
req.logTheSuccess();
|
||||
return req;
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////
|
||||
PLATFORM_IMPL(conv2d_bp, ENGINE_CPU) {
|
||||
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_NULLIFIED(0); // [bS, iH, iW, iC] (NHWC) or [bS, iC, iH, iW] (NCHW), epsilon
|
||||
auto gradW = OUTPUT_NULLIFIED(1); // [kH, kW, iC, oC], [oC, iC, kH, kW], [oC, kH, kW, iC]
|
||||
auto gradB = block.width() > 3 ? OUTPUT_NULLIFIED(2) : nullptr; // [oC]
|
||||
|
||||
sd::LongType kH = INT_ARG(0); // filter(kernel) height
|
||||
sd::LongType kW = INT_ARG(1); // filter(kernel) width
|
||||
sd::LongType sH = INT_ARG(2); // strides height
|
||||
sd::LongType sW = INT_ARG(3); // strides width
|
||||
sd::LongType pH = INT_ARG(4); // paddings height
|
||||
sd::LongType pW = INT_ARG(5); // paddings width
|
||||
sd::LongType dH = INT_ARG(6); // dilations height
|
||||
sd::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]
|
||||
|
||||
sd::LongType bS, iC, iH, iW, oC, oH,
|
||||
oW; // batch size, input channels, input height/width, output channels, output height/width;
|
||||
sd::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);
|
||||
|
||||
sd::LongType trueoH, trueoW; // true output height, width
|
||||
ConvolutionUtils::calcOutSizePool2D(trueoH, trueoW, kH, kW, sH, sW, pH, pW, dH, dW, iH, iW, paddingMode);
|
||||
|
||||
if (paddingMode) // SAME
|
||||
ConvolutionUtils::calcPadding2D(pH, pW, oH, oW, iH, iW, kH, kW, sH, sW, dH, dW, paddingMode);
|
||||
|
||||
std::vector<sd::LongType> expectedGradOShape =
|
||||
ShapeUtils::composeShapeUsingDimsAndIdx({bS, oC, trueoH, trueoW, 0, indIOioC, indOoH, indOoH + 1});
|
||||
std::vector<sd::LongType> expectedWeightsShape = ConvolutionUtils::expectWeightsShape(wFormat, kH, kW, iC, oC);
|
||||
REQUIRE_TRUE(
|
||||
gradO->isSameShape(expectedGradOShape), 0,
|
||||
"CONV2D_BP MKLDNN 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,
|
||||
"CONV2D_BP MKLDNN 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,
|
||||
"CONV2D_BP MKLDNN OP: wrong shape of array with biases, expected rank, length: <=2, %i, but got %i, "
|
||||
"%i instead !",
|
||||
oC, bias->rankOf(), bias->lengthOf());
|
||||
|
||||
conv2dBpMKLDNN(input, weights, bias, gradO, gradI, gradW, gradB, kH, kW, sH, sW, pH, pW, dH, dW, paddingMode, isNCHW,
|
||||
wFormat);
|
||||
|
||||
return sd::Status::OK;
|
||||
}
|
||||
|
||||
PLATFORM_CHECK(conv2d_bp, ENGINE_CPU) {
|
||||
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
|
||||
|
||||
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] always
|
||||
auto gradB = block.width() > 3 ? OUTPUT_VARIABLE(2) : nullptr; // [oC]
|
||||
|
||||
Requirements req("ONEDNN CONV2d_BP OP");
|
||||
req.expectTrue(block.isUseONEDNN(), IS_USE_ONEDNN_MSG) &&
|
||||
req.expectTrue(sd::ONEDNNStream::isSupported({input, weights, bias, gradO, gradI, gradW, gradB}),
|
||||
ONEDNN_STREAM_NOT_SUPPORTED);
|
||||
req.logTheSuccess();
|
||||
return req;
|
||||
}
|
||||
|
||||
} // namespace platforms
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,458 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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 saudet
|
||||
// @author raver119@gmail.com
|
||||
//
|
||||
#include <helpers/MKLDNNStream.h>
|
||||
#include <ops/declarable/OpRegistrator.h>
|
||||
#include <ops/declarable/PlatformHelper.h>
|
||||
#include <ops/declarable/helpers/convolutions.h>
|
||||
#include <system/platform_boilerplate.h>
|
||||
|
||||
#include "mkldnnUtils.h"
|
||||
|
||||
using namespace dnnl;
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace platforms {
|
||||
|
||||
//////////////////////////////////////////////////////////////////////
|
||||
static void conv3dMKLDNN(NDArray *input, NDArray *weights, NDArray *bias, NDArray *output,
|
||||
const sd::LongType kD, const sd::LongType kH, const sd::LongType kW, const sd::LongType sD, const sd::LongType sH, const sd::LongType sW,
|
||||
const sd::LongType pD, const sd::LongType pH, const sd::LongType pW, const sd::LongType dD, const sd::LongType dH, const sd::LongType dW,
|
||||
const int paddingMode, const int isNCDHW, const int wFormat) {
|
||||
// mkl support weights in [oC, iC, kD, kH, kW] format only
|
||||
|
||||
sd::LongType bS, iC, iD, iH, iW, oC, oD, oH,
|
||||
oW; // batch size, input channels, input depth/height/width, output channels, output depth/height/width;
|
||||
sd::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);
|
||||
|
||||
// const int pWSame = (paddingMode == 2 && dW > 1) ? ((oW - 1) * sW + (kW - 1) * dW + 1 - iW) / 2 : pW; // dH ==
|
||||
// 1 for causal mode in conv1d
|
||||
|
||||
dnnl::memory::dims strides = {sD, sH, sW};
|
||||
dnnl::memory::dims padding = {pD, pH, pW};
|
||||
// dnnl::memory::dims padding_r = { (oH - 1) * sH - iH + kH - pH, (oW - 1) * sW - iW + kW - pWSame };
|
||||
dnnl::memory::dims padding_r = {(oD - 1) * sD - iD + kD - pD, (oH - 1) * sH - iH + kH - pH,
|
||||
(oW - 1) * sW - iW + kW - pW};
|
||||
dnnl::memory::dims dilation = {dD - 1, dH - 1, dW - 1};
|
||||
|
||||
auto xzFormatMkl = isNCDHW ? dnnl::memory::format_tag::ncdhw : dnnl::memory::format_tag::ndhwc;
|
||||
dnnl::memory::format_tag wFormatMkl = dnnl::memory::format_tag::oidhw;
|
||||
|
||||
dnnl::memory::dims xDims = {bS, iC, iD, iH, iW};
|
||||
dnnl::memory::dims wDims = {oC, iC, kD, kH, kW};
|
||||
dnnl::memory::dims zDims = {bS, oC, oD, oH, oW};
|
||||
|
||||
std::vector<int> permut;
|
||||
if (0 == wFormat)
|
||||
permut = {4, 3, 0, 1, 2}; // [kD, kH, kW, iC, oC] -> [oC, iC, kD, kH, kW]
|
||||
else if (2 == wFormat)
|
||||
permut = {0, 4, 1, 2, 3}; // [oC, kD, kH, kW, iC] -> [oC, iC, kD, kH, kW]
|
||||
|
||||
auto type = dnnl::memory::data_type::f32;
|
||||
|
||||
// memory descriptors for arrays
|
||||
|
||||
// input
|
||||
dnnl::memory::desc x_mkl_md = dnnl::memory::desc(xDims, type, dnnl::memory::format_tag::any);
|
||||
dnnl::memory::desc x_user_md = dnnl::memory::desc(xDims, type, xzFormatMkl);
|
||||
onednnUtils::setBlockStrides(*input, x_user_md);
|
||||
|
||||
// weights
|
||||
dnnl::memory::desc w_mkl_md = dnnl::memory::desc(wDims, type, dnnl::memory::format_tag::any);
|
||||
dnnl::memory::desc w_user_md = dnnl::memory::desc(wDims, type, wFormatMkl);
|
||||
onednnUtils::setBlockStrides(*weights, w_user_md, permut);
|
||||
|
||||
// bias
|
||||
dnnl::memory::desc b_mkl_md;
|
||||
if (bias != nullptr) b_mkl_md = dnnl::memory::desc({oC}, type, dnnl::memory::format_tag::x);
|
||||
|
||||
// output
|
||||
dnnl::memory::desc z_mkl_md = dnnl::memory::desc(zDims, type, dnnl::memory::format_tag::any);
|
||||
dnnl::memory::desc z_user_md = dnnl::memory::desc(zDims, type, xzFormatMkl);
|
||||
onednnUtils::setBlockStrides(*output, z_user_md);
|
||||
|
||||
auto engine = onednnUtils::getEngine(LaunchContext::defaultContext()->engine());
|
||||
|
||||
// operation primitive description
|
||||
dnnl::convolution_forward::desc op_desc(dnnl::prop_kind::forward_inference, dnnl::algorithm::convolution_auto,
|
||||
x_mkl_md, w_mkl_md, b_mkl_md, z_mkl_md, strides, dilation, padding,
|
||||
padding_r);
|
||||
dnnl::convolution_forward::primitive_desc op_prim_desc(op_desc, engine);
|
||||
|
||||
// arguments (memory buffers) necessary for calculations
|
||||
std::unordered_map<int, dnnl::memory> args;
|
||||
|
||||
dnnl::stream stream(engine);
|
||||
|
||||
// provide memory buffers and check whether reorder is required
|
||||
|
||||
// input
|
||||
onednnUtils::loadDataToMklStream(*input, engine, stream, x_user_md, op_prim_desc.src_desc(), args[DNNL_ARG_SRC]);
|
||||
|
||||
// weights
|
||||
onednnUtils::loadDataToMklStream(*weights, engine, stream, w_user_md, op_prim_desc.weights_desc(),
|
||||
args[DNNL_ARG_WEIGHTS]);
|
||||
|
||||
// bias
|
||||
if (bias != nullptr) {
|
||||
auto b_mkl_mem = dnnl::memory(b_mkl_md, engine, const_cast<void *>(bias->buffer()));
|
||||
args[DNNL_ARG_BIAS] = b_mkl_mem;
|
||||
}
|
||||
|
||||
// output
|
||||
auto z_user_mem =
|
||||
onednnUtils::loadDataToMklStream(*output, engine, stream, z_user_md, op_prim_desc.dst_desc(), args[DNNL_ARG_DST]);
|
||||
|
||||
// run calculations
|
||||
dnnl::convolution_forward(op_prim_desc).execute(stream, args);
|
||||
|
||||
// reorder outputs if necessary
|
||||
if (op_prim_desc.dst_desc() != z_user_mem.get_desc())
|
||||
dnnl::reorder(args[DNNL_ARG_DST], z_user_mem).execute(stream, args[DNNL_ARG_DST], z_user_mem);
|
||||
|
||||
stream.wait();
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////
|
||||
static void conv3dBpMKLDNN(NDArray *input, NDArray *weights, NDArray *bias, NDArray *gradO,
|
||||
NDArray *gradI, NDArray *gradW, NDArray *gradB, const sd::LongType kD, const sd::LongType kH, const sd::LongType kW,
|
||||
const sd::LongType sD, const sd::LongType sH, const sd::LongType sW, const sd::LongType pD, const sd::LongType pH, const sd::LongType pW,
|
||||
const sd::LongType dD, const sd::LongType dH, const sd::LongType dW, const int paddingMode, const int isNCDHW,
|
||||
const int wFormat) {
|
||||
// mkl support weights/gradW in [oC, iC, kD, kH, kW] format only
|
||||
|
||||
sd::LongType bS, iC, iD, iH, iW, oC, oD, oH,
|
||||
oW; // batch size, input channels, input depth/height/width, output channels, output depth/height/width;
|
||||
sd::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);
|
||||
|
||||
// const int pWSame = (paddingMode == 2 && dW > 1) ? ((oW - 1) * sW + (kW - 1) * dW + 1 - iW) / 2 : pW; // dH ==
|
||||
// 1 for causal mode in conv1d
|
||||
|
||||
dnnl::memory::dims strides = {sD, sH, sW};
|
||||
dnnl::memory::dims padding = {pD, pH, pW};
|
||||
// dnnl::memory::dims padding_r = { (oH - 1) * sH - iH + kH - pH, (oW - 1) * sW - iW + kW - pWSame };
|
||||
dnnl::memory::dims padding_r = {(oD - 1) * sD - iD + kD - pD, (oH - 1) * sH - iH + kH - pH,
|
||||
(oW - 1) * sW - iW + kW - pW};
|
||||
dnnl::memory::dims dilation = {dD - 1, dH - 1, dW - 1};
|
||||
|
||||
auto xzFormatMkl = isNCDHW ? dnnl::memory::format_tag::ncdhw : dnnl::memory::format_tag::ndhwc;
|
||||
dnnl::memory::format_tag wFormatMkl = dnnl::memory::format_tag::oidhw;
|
||||
|
||||
dnnl::memory::dims xDims = {bS, iC, iD, iH, iW};
|
||||
dnnl::memory::dims wDims = {oC, iC, kD, kH, kW};
|
||||
dnnl::memory::dims zDims = {bS, oC, oD, oH, oW};
|
||||
|
||||
auto type = dnnl::memory::data_type::f32;
|
||||
|
||||
std::vector<int> permut;
|
||||
if (0 == wFormat)
|
||||
permut = {4, 3, 0, 1, 2}; // [kD, kH, kW, iC, oC] -> [oC, iC, kD, kH, kW]
|
||||
else if (2 == wFormat)
|
||||
permut = {0, 4, 1, 2, 3}; // [oC, kD, kH, kW, iC] -> [oC, iC, kD, kH, kW]
|
||||
|
||||
// memory descriptors for arrays
|
||||
|
||||
// input
|
||||
dnnl::memory::desc x_mkl_md = dnnl::memory::desc(xDims, type, dnnl::memory::format_tag::any);
|
||||
dnnl::memory::desc x_user_md = dnnl::memory::desc(xDims, type, xzFormatMkl);
|
||||
onednnUtils::setBlockStrides(*input, x_user_md);
|
||||
|
||||
// weights
|
||||
dnnl::memory::desc w_mkl_md = dnnl::memory::desc(wDims, type, dnnl::memory::format_tag::any);
|
||||
dnnl::memory::desc w_user_md = dnnl::memory::desc(wDims, type, wFormatMkl);
|
||||
onednnUtils::setBlockStrides(*weights, w_user_md, permut);
|
||||
|
||||
// gradO
|
||||
dnnl::memory::desc gradO_mkl_md = dnnl::memory::desc(zDims, type, dnnl::memory::format_tag::any);
|
||||
dnnl::memory::desc gradO_user_md = dnnl::memory::desc(zDims, type, xzFormatMkl);
|
||||
|
||||
onednnUtils::setBlockStrides(*gradO, gradO_user_md);
|
||||
|
||||
// gradI
|
||||
dnnl::memory::desc gradI_mkl_md = dnnl::memory::desc(xDims, type, dnnl::memory::format_tag::any);
|
||||
dnnl::memory::desc gradI_user_md = dnnl::memory::desc(xDims, type, xzFormatMkl);
|
||||
|
||||
onednnUtils::setBlockStrides(*gradI, gradI_user_md);
|
||||
|
||||
// gradW
|
||||
dnnl::memory::desc gradW_mkl_md = dnnl::memory::desc(wDims, type, dnnl::memory::format_tag::any);
|
||||
dnnl::memory::desc gradW_user_md = dnnl::memory::desc(wDims, type, wFormatMkl);
|
||||
onednnUtils::setBlockStrides(*gradW, gradW_user_md, permut);
|
||||
|
||||
// gradB
|
||||
dnnl::memory::desc gradB_mkl_md;
|
||||
if (gradB != nullptr) gradB_mkl_md = dnnl::memory::desc({oC}, type, dnnl::memory::format_tag::x);
|
||||
|
||||
auto engine = onednnUtils::getEngine(LaunchContext::defaultContext()->engine());
|
||||
|
||||
// forward primitive description
|
||||
dnnl::convolution_forward::desc op_ff_desc(dnnl::prop_kind::forward_inference, dnnl::algorithm::convolution_auto,
|
||||
x_mkl_md, w_mkl_md, gradB_mkl_md, gradO_mkl_md, strides, dilation, padding,
|
||||
padding_r);
|
||||
dnnl::convolution_forward::primitive_desc op_ff_prim_desc(op_ff_desc, engine);
|
||||
|
||||
// backward data primitive description
|
||||
dnnl::convolution_backward_data::desc op_data_bp_desc(dnnl::algorithm::convolution_auto, gradI_mkl_md, w_mkl_md,
|
||||
gradO_mkl_md, strides, dilation, padding, padding_r);
|
||||
dnnl::convolution_backward_data::primitive_desc op_data_bp_prim_desc(op_data_bp_desc, engine, op_ff_prim_desc);
|
||||
|
||||
// backward weights primitive description
|
||||
dnnl::convolution_backward_weights::desc op_weights_bp_desc(dnnl::algorithm::convolution_auto, x_mkl_md, gradW_mkl_md,
|
||||
gradB_mkl_md, gradO_mkl_md, strides, dilation, padding,
|
||||
padding_r);
|
||||
dnnl::convolution_backward_weights::primitive_desc op_weights_bp_prim_desc(op_weights_bp_desc, engine,
|
||||
op_ff_prim_desc);
|
||||
|
||||
// arguments (memory buffers) necessary for calculations
|
||||
std::unordered_map<int, dnnl::memory> args;
|
||||
|
||||
dnnl::stream stream(engine);
|
||||
|
||||
// provide memory buffers and check whether reorder is required
|
||||
|
||||
// input
|
||||
onednnUtils::loadDataToMklStream(*input, engine, stream, x_user_md, op_weights_bp_prim_desc.src_desc(),
|
||||
args[DNNL_ARG_SRC]);
|
||||
|
||||
// weights
|
||||
onednnUtils::loadDataToMklStream(*weights, engine, stream, w_user_md, op_data_bp_prim_desc.weights_desc(),
|
||||
args[DNNL_ARG_WEIGHTS]);
|
||||
|
||||
// gradO
|
||||
auto gradO_user_mem = dnnl::memory(gradO_user_md, engine, const_cast<void *>(gradO->buffer()));
|
||||
const bool gradOReorderW = op_weights_bp_prim_desc.diff_dst_desc() != gradO_user_mem.get_desc();
|
||||
const bool gradOReorderD = op_data_bp_prim_desc.diff_dst_desc() != gradO_user_mem.get_desc();
|
||||
auto gradO_mkl_memW = gradOReorderW ? dnnl::memory(op_weights_bp_prim_desc.diff_dst_desc(), engine) : gradO_user_mem;
|
||||
auto gradO_mkl_memD = gradOReorderD ? dnnl::memory(op_data_bp_prim_desc.diff_dst_desc(), engine) : gradO_user_mem;
|
||||
if (gradOReorderW) dnnl::reorder(gradO_user_mem, gradO_mkl_memW).execute(stream, gradO_user_mem, gradO_mkl_memW);
|
||||
if (gradOReorderD) dnnl::reorder(gradO_user_mem, gradO_mkl_memD).execute(stream, gradO_user_mem, gradO_mkl_memD);
|
||||
args[DNNL_ARG_DIFF_DST] = gradO_mkl_memD;
|
||||
|
||||
// gradI
|
||||
auto gradI_user_mem = onednnUtils::loadDataToMklStream(*gradI, engine, stream, gradI_user_md,
|
||||
op_data_bp_prim_desc.diff_src_desc(), args[DNNL_ARG_DIFF_SRC]);
|
||||
|
||||
// gradW
|
||||
auto gradW_user_mem = onednnUtils::loadDataToMklStream(
|
||||
*gradW, engine, stream, gradW_user_md, op_weights_bp_prim_desc.diff_weights_desc(), args[DNNL_ARG_DIFF_WEIGHTS]);
|
||||
|
||||
// gradB
|
||||
if (gradB != nullptr) {
|
||||
auto gradB_mkl_mem = dnnl::memory(gradB_mkl_md, engine, gradB->buffer());
|
||||
args[DNNL_ARG_DIFF_BIAS] = gradB_mkl_mem;
|
||||
}
|
||||
|
||||
// run backward data calculations
|
||||
dnnl::convolution_backward_data(op_data_bp_prim_desc).execute(stream, args);
|
||||
|
||||
if (gradOReorderW || gradOReorderD) args[DNNL_ARG_DIFF_DST] = gradO_mkl_memW;
|
||||
|
||||
// run backward weights calculations
|
||||
dnnl::convolution_backward_weights(op_weights_bp_prim_desc).execute(stream, args);
|
||||
|
||||
// reorder gradI if necessary
|
||||
if (op_data_bp_prim_desc.diff_src_desc() != gradI_user_mem.get_desc())
|
||||
dnnl::reorder(args[DNNL_ARG_DIFF_SRC], gradI_user_mem).execute(stream, args[DNNL_ARG_DIFF_SRC], gradI_user_mem);
|
||||
if (op_weights_bp_prim_desc.diff_weights_desc() != gradW_user_mem.get_desc())
|
||||
dnnl::reorder(args[DNNL_ARG_DIFF_WEIGHTS], gradW_user_mem)
|
||||
.execute(stream, args[DNNL_ARG_DIFF_WEIGHTS], gradW_user_mem);
|
||||
|
||||
stream.wait();
|
||||
|
||||
}
|
||||
|
||||
|
||||
|
||||
//////////////////////////////////////////////////////////////////////
|
||||
PLATFORM_IMPL(conv3dnew, ENGINE_CPU) {
|
||||
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,
|
||||
"CUSTOM CONV3D MKLDNN OP: rank of input array must be equal to 5, but got %i instead !",
|
||||
input->rankOf());
|
||||
REQUIRE_TRUE(weights->rankOf() == 5, 0,
|
||||
"CUSTOM CONV3D MKLDNN OP: rank of weights array must be equal to 5, but got %i instead !",
|
||||
weights->rankOf());
|
||||
|
||||
sd::LongType kD = INT_ARG(0) > 0 ? INT_ARG(0) : static_cast<sd::LongType>(weights->sizeAt(0)); // filter(kernel) depth
|
||||
sd::LongType kH = INT_ARG(1) > 0 ? INT_ARG(1) : static_cast<sd::LongType>(weights->sizeAt(1)); // filter(kernel) height
|
||||
sd::LongType kW = INT_ARG(2) > 0 ? INT_ARG(2) : static_cast<sd::LongType>(weights->sizeAt(2)); // filter(kernel) width
|
||||
sd::LongType sD = INT_ARG(3); // strides depth
|
||||
sd::LongType sH = INT_ARG(4); // strides height
|
||||
sd::LongType sW = INT_ARG(5); // strides width
|
||||
sd::LongType pD = INT_ARG(6); // paddings depth
|
||||
sd::LongType pH = INT_ARG(7); // paddings height
|
||||
sd::LongType pW = INT_ARG(8); // paddings width
|
||||
sd::LongType dD = INT_ARG(9); // dilations depth
|
||||
sd::LongType dH = INT_ARG(10); // dilations height
|
||||
sd::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]
|
||||
|
||||
sd::LongType bS, iC, iD, iH, iW, oC, oD, oH,
|
||||
oW; // batch size, input channels, input depth/height/width, output channels, output depth/height/width;
|
||||
sd::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);
|
||||
|
||||
std::vector<sd::LongType> expectedWeightsShape = ConvolutionUtils::expectWeightsShape(wFormat, kD, kH, kW, iC, oC);
|
||||
REQUIRE_TRUE(weights->isSameShape(expectedWeightsShape), 0,
|
||||
"CUSTOM CONV3D MKLDNN 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 CONV3D MKLDNN OP: wrong shape of array with biases, expected rank, length: <=2, %i, but got "
|
||||
"%i, %i instead !",
|
||||
oC, bias->rankOf(), bias->lengthOf());
|
||||
|
||||
if (paddingMode) // SAME
|
||||
ConvolutionUtils::calcPadding3D(pD, pH, pW, oD, oH, oW, iD, iH, iW, kD, kH, kW, sD, sH, sW, dD, dH, dW);
|
||||
|
||||
conv3dMKLDNN(input, weights, bias, output, kD, kH, kW, sD, sH, sW, pD, pH, pW, dD, dH, dW, paddingMode, isNCDHW,
|
||||
wFormat);
|
||||
|
||||
return sd::Status::OK;
|
||||
}
|
||||
|
||||
PLATFORM_CHECK(conv3dnew, ENGINE_CPU) {
|
||||
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] always
|
||||
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)
|
||||
Requirements req("ONEDNN CONV3d OP");
|
||||
req.expectTrue(block.isUseONEDNN(), IS_USE_ONEDNN_MSG) &&
|
||||
req.expectTrue(sd::ONEDNNStream::isSupported({input, weights, bias, output}), ONEDNN_STREAM_NOT_SUPPORTED);
|
||||
req.logTheSuccess();
|
||||
return req;
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////
|
||||
PLATFORM_IMPL(conv3dnew_bp, ENGINE_CPU) {
|
||||
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_NULLIFIED(0); // [bS, iD, iH, iW, iC] (NDHWC) or [bS, iC, iD, iH, iW] (NCDHW), epsilon
|
||||
auto gradW = OUTPUT_NULLIFIED(1); // [kD, kH, kW, iC, oC], [oC, iC, kD, kH, kW], [oC, kD, kH, kW, iC]
|
||||
auto gradB = block.width() > 3 ? OUTPUT_NULLIFIED(2) : nullptr; // [oC]
|
||||
|
||||
REQUIRE_TRUE(input->rankOf() == 5, 0,
|
||||
"CUSTOM CONV3D_BP MKLDNN OP: rank of input array must be equal to 5, but got %i instead !",
|
||||
input->rankOf());
|
||||
REQUIRE_TRUE(weights->rankOf() == 5, 0,
|
||||
"CUSTOM CONV3D_BP MKLDNN OP: rank of weights array must be equal to 5, but got %i instead !",
|
||||
weights->rankOf());
|
||||
REQUIRE_TRUE(gradO->rankOf() == 5, 0,
|
||||
"CUSTOM CONV3D_BP MKLDNN OP: rank of output gradients (next epsilon) array must be equal to 5, but got "
|
||||
"%i instead !",
|
||||
gradO->rankOf());
|
||||
|
||||
sd::LongType kD = INT_ARG(0) > 0 ? INT_ARG(0) : static_cast<sd::LongType>(weights->sizeAt(0)); // filter(kernel) depth
|
||||
sd::LongType kH = INT_ARG(1) > 0 ? INT_ARG(1) : static_cast<sd::LongType>(weights->sizeAt(1)); // filter(kernel) height
|
||||
sd::LongType kW = INT_ARG(2) > 0 ? INT_ARG(2) : static_cast<sd::LongType>(weights->sizeAt(2)); // filter(kernel) width
|
||||
sd::LongType sD = INT_ARG(3); // strides depth
|
||||
sd::LongType sH = INT_ARG(4); // strides height
|
||||
sd::LongType sW = INT_ARG(5); // strides width
|
||||
sd::LongType pD = INT_ARG(6); // paddings depth
|
||||
sd::LongType pH = INT_ARG(7); // paddings height
|
||||
sd::LongType pW = INT_ARG(8); // paddings width
|
||||
sd::LongType dD = INT_ARG(9); // dilations depth
|
||||
sd::LongType dH = INT_ARG(10); // dilations height
|
||||
sd::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]
|
||||
|
||||
sd::LongType bS, iC, iD, iH, iW, oC, oD, oH,
|
||||
oW; // batch size, input channels, input depth/height/width, output channels, output depth/height/width;
|
||||
sd::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);
|
||||
|
||||
if (paddingMode) // SAME
|
||||
ConvolutionUtils::calcPadding3D(pD, pH, pW, oD, oH, oW, iD, iH, iW, kD, kH, kW, sD, sH, sW, dD, dH, dW);
|
||||
|
||||
sd::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);
|
||||
|
||||
std::vector<sd::LongType> expectedGradOShape = ShapeUtils::composeShapeUsingDimsAndIdx(
|
||||
{bS, oC, trueoD, trueoH, trueoW, 0, indIOioC, indIOioD, indIOioD + 1, indIOioD + 2});
|
||||
std::vector<sd::LongType> expectedWeightsShape = ConvolutionUtils::expectWeightsShape(wFormat, kD, kH, kW, iC, oC);
|
||||
REQUIRE_TRUE(
|
||||
gradO->isSameShape(expectedGradOShape), 0,
|
||||
"CUSTOM CONV3D_BP 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 CONV3D_BP 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 CONV3D_BP OP: wrong shape of array with biases, expected rank, length: <=2, %i, but got %i, "
|
||||
"%i instead !",
|
||||
oC, bias->rankOf(), bias->lengthOf());
|
||||
|
||||
conv3dBpMKLDNN(input, weights, bias, gradO, gradI, gradW, gradB, kD, kH, kW, sD, sH, sW, pD, pH, pW, dD, dH, dW,
|
||||
paddingMode, isNCDHW, wFormat);
|
||||
|
||||
return sd::Status::OK;
|
||||
}
|
||||
|
||||
PLATFORM_CHECK(conv3dnew_bp, ENGINE_CPU) {
|
||||
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]
|
||||
|
||||
Requirements req("ONEDNN CONV3d_BP OP");
|
||||
req.expectTrue(block.isUseONEDNN(), IS_USE_ONEDNN_MSG) &&
|
||||
req.expectTrue(sd::ONEDNNStream::isSupported({input, weights, bias, gradO, gradI, gradW, gradB}),
|
||||
ONEDNN_STREAM_NOT_SUPPORTED);
|
||||
req.logTheSuccess();
|
||||
return req;
|
||||
}
|
||||
|
||||
} // namespace platforms
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,532 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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 <helpers/MKLDNNStream.h>
|
||||
#include <ops/declarable/OpRegistrator.h>
|
||||
#include <ops/declarable/PlatformHelper.h>
|
||||
#include <ops/declarable/helpers/convolutions.h>
|
||||
#include <system/platform_boilerplate.h>
|
||||
|
||||
#include "mkldnnUtils.h"
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace platforms {
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
static void deconv2dMKLDNN(NDArray* input, NDArray* weights, NDArray* bias, NDArray* output,
|
||||
const sd::LongType kH, const sd::LongType kW, const sd::LongType sH, const sd::LongType sW, const sd::LongType pH, const sd::LongType pW,
|
||||
const sd::LongType dH, const sd::LongType dW, const int paddingMode, const bool isNCHW, const int wFormat) {
|
||||
// mkl supports weights format [oC, iC, kH, kW] only
|
||||
|
||||
sd::LongType bS, iC, iH, iW, oC, oH,
|
||||
oW; // batch size, input channels, input height/width, output channels, output height/width;
|
||||
sd::LongType indIOioC, indIiH, indWoC, indWiC, indWkH, indOoH; // corresponding indexes
|
||||
ConvolutionUtils::getSizesAndIndexesConv2d(isNCHW, wFormat, *input, *output, bS, iC, iH, iW, oC, oH, oW, indIOioC,
|
||||
indIiH, indWoC, indWiC, indWkH, indOoH);
|
||||
|
||||
dnnl::memory::dims strides = {sH, sW};
|
||||
dnnl::memory::dims padding = {pH, pW};
|
||||
dnnl::memory::dims padding_r = {(iH - 1) * sH - oH + kH - pH, (iW - 1) * sW - oW + kW - pW};
|
||||
dnnl::memory::dims dilation = {dH - 1, dW - 1};
|
||||
|
||||
std::vector<int> permut;
|
||||
if (0 == wFormat)
|
||||
permut = {2, 3, 0, 1}; // [kH, kW, oC, iC] -> [oC, iC, kH, kW]
|
||||
else if (1 == wFormat)
|
||||
permut = {1, 0, 2, 3}; // [iC, oC, kH, kW] -> [oC, iC, kH, kW]
|
||||
else
|
||||
permut = {3, 0, 1, 2}; // [iC, kH, kW, oC] -> [oC, iC, kH, kW]
|
||||
|
||||
// input type
|
||||
dnnl::memory::data_type xType;
|
||||
if (input->dataType() == DataType::FLOAT32)
|
||||
xType = dnnl::memory::data_type::f32;
|
||||
else if (input->dataType() == DataType::HALF)
|
||||
xType = dnnl::memory::data_type::f16;
|
||||
else if (input->dataType() == DataType::UINT8)
|
||||
xType = dnnl::memory::data_type::u8;
|
||||
else
|
||||
xType = dnnl::memory::data_type::s8;
|
||||
|
||||
// weights type
|
||||
dnnl::memory::data_type wType = xType;
|
||||
if (xType == dnnl::memory::data_type::u8) wType = dnnl::memory::data_type::s8;
|
||||
|
||||
// output and bias type (have the same types)
|
||||
dnnl::memory::data_type zType;
|
||||
if (output->dataType() == DataType::FLOAT32)
|
||||
zType = dnnl::memory::data_type::f32;
|
||||
else if (output->dataType() == DataType::HALF)
|
||||
zType = dnnl::memory::data_type::f16;
|
||||
else if (output->dataType() == DataType::UINT8)
|
||||
zType = dnnl::memory::data_type::u8;
|
||||
else if (output->dataType() == DataType::INT8)
|
||||
zType = dnnl::memory::data_type::s8;
|
||||
else
|
||||
zType = dnnl::memory::data_type::s32;
|
||||
|
||||
dnnl::memory::format_tag xFormatMkl = isNCHW ? dnnl::memory::format_tag::nchw : dnnl::memory::format_tag::nhwc;
|
||||
dnnl::memory::format_tag wFormatMkl = dnnl::memory::format_tag::oihw;
|
||||
|
||||
dnnl::memory::dims xDims = {bS, iC, iH, iW};
|
||||
dnnl::memory::dims wDims = {oC, iC, kH, kW};
|
||||
dnnl::memory::dims zDims = {bS, oC, oH, oW};
|
||||
|
||||
// memory descriptors for arrays
|
||||
|
||||
// input
|
||||
dnnl::memory::desc x_mkl_md = dnnl::memory::desc(xDims, xType, dnnl::memory::format_tag::any);
|
||||
dnnl::memory::desc x_user_md = dnnl::memory::desc(xDims, xType, xFormatMkl);
|
||||
onednnUtils::setBlockStrides(*input, x_user_md);
|
||||
|
||||
// weights
|
||||
dnnl::memory::desc w_mkl_md = dnnl::memory::desc(wDims, wType, dnnl::memory::format_tag::any);
|
||||
dnnl::memory::desc w_user_md = dnnl::memory::desc(wDims, wType, wFormatMkl);
|
||||
onednnUtils::setBlockStrides(*weights, w_user_md, permut);
|
||||
|
||||
// bias
|
||||
dnnl::memory::desc b_mkl_md;
|
||||
if (bias != nullptr) b_mkl_md = dnnl::memory::desc({oC}, zType, dnnl::memory::format_tag::x);
|
||||
|
||||
// output
|
||||
dnnl::memory::desc z_mkl_md = dnnl::memory::desc(zDims, zType, dnnl::memory::format_tag::any);
|
||||
dnnl::memory::desc z_user_md = dnnl::memory::desc(zDims, zType, xFormatMkl);
|
||||
onednnUtils::setBlockStrides(*output, z_user_md);
|
||||
|
||||
auto engine = onednnUtils::getEngine(LaunchContext::defaultContext()->engine());
|
||||
|
||||
// operation primitive description
|
||||
dnnl::deconvolution_forward::desc op_desc(dnnl::prop_kind::forward_inference, dnnl::algorithm::deconvolution_direct,
|
||||
x_mkl_md, w_mkl_md, b_mkl_md, z_mkl_md, strides, dilation, padding,
|
||||
padding_r);
|
||||
dnnl::deconvolution_forward::primitive_desc op_prim_desc(op_desc, engine);
|
||||
|
||||
// arguments (memory buffers) necessary for calculations
|
||||
std::unordered_map<int, dnnl::memory> args;
|
||||
|
||||
dnnl::stream stream(engine);
|
||||
|
||||
// provide memory buffers and check whether reorder is required
|
||||
|
||||
// input
|
||||
onednnUtils::loadDataToMklStream(*input, engine, stream, x_user_md, op_prim_desc.src_desc(), args[DNNL_ARG_SRC]);
|
||||
|
||||
// weights
|
||||
onednnUtils::loadDataToMklStream(*weights, engine, stream, w_user_md, op_prim_desc.weights_desc(),
|
||||
args[DNNL_ARG_WEIGHTS]);
|
||||
|
||||
// bias
|
||||
if (bias != nullptr) {
|
||||
auto b_mkl_mem = dnnl::memory(b_mkl_md, engine, const_cast<void*>(bias->buffer()));
|
||||
args[DNNL_ARG_BIAS] = b_mkl_mem;
|
||||
}
|
||||
|
||||
// output
|
||||
auto z_user_mem =
|
||||
onednnUtils::loadDataToMklStream(*output, engine, stream, z_user_md, op_prim_desc.dst_desc(), args[DNNL_ARG_DST]);
|
||||
|
||||
// run calculations
|
||||
dnnl::deconvolution_forward(op_prim_desc).execute(stream, args);
|
||||
|
||||
// reorder outputs if necessary
|
||||
if (op_prim_desc.dst_desc() != z_user_mem.get_desc())
|
||||
dnnl::reorder(args[DNNL_ARG_DST], z_user_mem).execute(stream, args[DNNL_ARG_DST], z_user_mem);
|
||||
|
||||
stream.wait();
|
||||
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
static void deconv2dBpMKLDNN(NDArray* input, NDArray* weights, NDArray* gradO, NDArray* gradI,
|
||||
NDArray* gradW, NDArray* gradB, const sd::LongType kH, const sd::LongType kW, const sd::LongType sH, const sd::LongType sW,
|
||||
const sd::LongType pH, const sd::LongType pW, const sd::LongType dH, const sd::LongType dW, const int paddingMode,
|
||||
const bool isNCHW, const int wFormat) {
|
||||
// mkl supports weights/gradW in [oC, iC, kH, kW] format only
|
||||
|
||||
sd::LongType bS, iC, iH, iW, oC, oH,
|
||||
oW; // batch size, input channels, input height/width, output channels, output height/width;
|
||||
sd::LongType indIOioC, indIiH, indWoC, indWiC, indWkH, indOoH; // corresponding indexes
|
||||
ConvolutionUtils::getSizesAndIndexesConv2d(isNCHW, wFormat, *input, *gradO, bS, iC, iH, iW, oC, oH, oW, indIOioC,
|
||||
indIiH, indWoC, indWiC, indWkH, indOoH);
|
||||
|
||||
dnnl::memory::dims strides = {sH, sW};
|
||||
dnnl::memory::dims padding = {pH, pW};
|
||||
dnnl::memory::dims padding_r = {(iH - 1) * sH - oH + kH - pH, (iW - 1) * sW - oW + kW - pW};
|
||||
dnnl::memory::dims dilation = {dH - 1, dW - 1};
|
||||
|
||||
std::vector<int> permut;
|
||||
if (0 == wFormat)
|
||||
permut = {2, 3, 0, 1}; // [kH, kW, oC, iC] -> [oC, iC, kH, kW]
|
||||
else if (1 == wFormat)
|
||||
permut = {1, 0, 2, 3}; // [iC, oC, kH, kW] -> [oC, iC, kH, kW]
|
||||
else
|
||||
permut = {3, 0, 1, 2}; // [iC, kH, kW, oC] -> [oC, iC, kH, kW]
|
||||
|
||||
// input type
|
||||
dnnl::memory::data_type xType =
|
||||
input->dataType() == DataType::FLOAT32 ? dnnl::memory::data_type::f32 : dnnl::memory::data_type::bf16;
|
||||
// weights type
|
||||
dnnl::memory::data_type wType =
|
||||
weights->dataType() == DataType::FLOAT32 ? dnnl::memory::data_type::f32 : dnnl::memory::data_type::bf16;
|
||||
// gradO type
|
||||
dnnl::memory::data_type gradOType =
|
||||
gradO->dataType() == DataType::FLOAT32 ? dnnl::memory::data_type::f32 : dnnl::memory::data_type::bf16;
|
||||
// gradI type
|
||||
dnnl::memory::data_type gradIType =
|
||||
gradI->dataType() == DataType::FLOAT32 ? dnnl::memory::data_type::f32 : dnnl::memory::data_type::bf16;
|
||||
// gradW type
|
||||
dnnl::memory::data_type gradWType =
|
||||
gradW->dataType() == DataType::FLOAT32 ? dnnl::memory::data_type::f32 : dnnl::memory::data_type::bf16;
|
||||
// gradB type
|
||||
dnnl::memory::data_type gradBType =
|
||||
gradB != nullptr
|
||||
? (gradB->dataType() == DataType::FLOAT32 ? dnnl::memory::data_type::f32 : dnnl::memory::data_type::bf16)
|
||||
: dnnl::memory::data_type::f32;
|
||||
|
||||
dnnl::memory::format_tag xFormatMkl = isNCHW ? dnnl::memory::format_tag::nchw : dnnl::memory::format_tag::nhwc;
|
||||
dnnl::memory::format_tag wFormatMkl = dnnl::memory::format_tag::oihw;
|
||||
|
||||
dnnl::memory::dims xDims = {bS, iC, iH, iW};
|
||||
dnnl::memory::dims wDims = {oC, iC, kH, kW};
|
||||
dnnl::memory::dims zDims = {bS, oC, oH, oW};
|
||||
|
||||
// memory descriptors for arrays
|
||||
|
||||
// input
|
||||
dnnl::memory::desc x_mkl_md = dnnl::memory::desc(xDims, xType, dnnl::memory::format_tag::any);
|
||||
dnnl::memory::desc x_user_md = dnnl::memory::desc(xDims, xType, xFormatMkl);
|
||||
onednnUtils::setBlockStrides(*input, x_user_md);
|
||||
|
||||
// weights
|
||||
dnnl::memory::desc w_mkl_md = dnnl::memory::desc(wDims, wType, dnnl::memory::format_tag::any);
|
||||
dnnl::memory::desc w_user_md = dnnl::memory::desc(wDims, wType, wFormatMkl);
|
||||
onednnUtils::setBlockStrides(*weights, w_user_md, permut);
|
||||
|
||||
// gradO
|
||||
dnnl::memory::desc gradO_mkl_md = dnnl::memory::desc(zDims, gradOType, dnnl::memory::format_tag::any);
|
||||
dnnl::memory::desc gradO_user_md = dnnl::memory::desc(zDims, gradOType, xFormatMkl);
|
||||
onednnUtils::setBlockStrides(*gradO, gradO_user_md);
|
||||
|
||||
// gradI
|
||||
dnnl::memory::desc gradI_mkl_md = dnnl::memory::desc(xDims, gradIType, dnnl::memory::format_tag::any);
|
||||
dnnl::memory::desc gradI_user_md = dnnl::memory::desc(xDims, gradIType, xFormatMkl);
|
||||
onednnUtils::setBlockStrides(*gradI, gradI_user_md);
|
||||
|
||||
// gradW
|
||||
dnnl::memory::desc gradW_mkl_md = dnnl::memory::desc(wDims, gradWType, dnnl::memory::format_tag::any);
|
||||
dnnl::memory::desc gradW_user_md = dnnl::memory::desc(wDims, gradWType, wFormatMkl);
|
||||
onednnUtils::setBlockStrides(*gradW, gradW_user_md, permut);
|
||||
|
||||
// gradB
|
||||
dnnl::memory::desc gradB_mkl_md;
|
||||
if (gradB != nullptr) gradB_mkl_md = dnnl::memory::desc({oC}, gradBType, dnnl::memory::format_tag::x);
|
||||
|
||||
auto engine = onednnUtils::getEngine(LaunchContext::defaultContext()->engine());
|
||||
|
||||
// forward primitive description
|
||||
dnnl::deconvolution_forward::desc op_ff_desc(dnnl::prop_kind::forward_inference,
|
||||
dnnl::algorithm::deconvolution_direct, x_mkl_md, w_mkl_md, gradB_mkl_md,
|
||||
gradO_mkl_md, strides, dilation, padding, padding_r);
|
||||
dnnl::deconvolution_forward::primitive_desc op_ff_prim_desc(op_ff_desc, engine);
|
||||
|
||||
// backward data primitive description
|
||||
dnnl::deconvolution_backward_data::desc op_data_bp_desc(dnnl::algorithm::deconvolution_direct, gradI_mkl_md, w_mkl_md,
|
||||
gradO_mkl_md, strides, dilation, padding, padding_r);
|
||||
dnnl::deconvolution_backward_data::primitive_desc op_data_bp_prim_desc(op_data_bp_desc, engine, op_ff_prim_desc);
|
||||
|
||||
// backward weights primitive description
|
||||
dnnl::deconvolution_backward_weights::desc op_weights_bp_desc(dnnl::algorithm::deconvolution_direct, x_mkl_md,
|
||||
gradW_mkl_md, gradB_mkl_md, gradO_mkl_md, strides,
|
||||
dilation, padding, padding_r);
|
||||
dnnl::deconvolution_backward_weights::primitive_desc op_weights_bp_prim_desc(op_weights_bp_desc, engine,
|
||||
op_ff_prim_desc);
|
||||
|
||||
// arguments (memory buffers) necessary for calculations
|
||||
std::unordered_map<int, dnnl::memory> args;
|
||||
|
||||
dnnl::stream stream(engine);
|
||||
|
||||
// provide memory buffers and check whether reorder is required
|
||||
|
||||
// input
|
||||
onednnUtils::loadDataToMklStream(*input, engine, stream, x_user_md, op_weights_bp_prim_desc.src_desc(),
|
||||
args[DNNL_ARG_SRC]);
|
||||
|
||||
// weights
|
||||
onednnUtils::loadDataToMklStream(*weights, engine, stream, w_user_md, op_data_bp_prim_desc.weights_desc(),
|
||||
args[DNNL_ARG_WEIGHTS]);
|
||||
|
||||
// gradO
|
||||
auto gradO_user_mem = dnnl::memory(gradO_user_md, engine, const_cast<void*>(gradO->buffer()));
|
||||
const bool gradOReorderW = op_weights_bp_prim_desc.diff_dst_desc() != gradO_user_mem.get_desc();
|
||||
const bool gradOReorderD = op_data_bp_prim_desc.diff_dst_desc() != gradO_user_mem.get_desc();
|
||||
auto gradO_mkl_memW = gradOReorderW ? dnnl::memory(op_weights_bp_prim_desc.diff_dst_desc(), engine) : gradO_user_mem;
|
||||
auto gradO_mkl_memD = gradOReorderD ? dnnl::memory(op_data_bp_prim_desc.diff_dst_desc(), engine) : gradO_user_mem;
|
||||
if (gradOReorderW) dnnl::reorder(gradO_user_mem, gradO_mkl_memW).execute(stream, gradO_user_mem, gradO_mkl_memW);
|
||||
if (gradOReorderD) dnnl::reorder(gradO_user_mem, gradO_mkl_memD).execute(stream, gradO_user_mem, gradO_mkl_memD);
|
||||
args[DNNL_ARG_DIFF_DST] = gradO_mkl_memD;
|
||||
|
||||
// gradI
|
||||
auto gradI_user_mem = onednnUtils::loadDataToMklStream(*gradI, engine, stream, gradI_user_md,
|
||||
op_data_bp_prim_desc.diff_src_desc(), args[DNNL_ARG_DIFF_SRC]);
|
||||
|
||||
// gradW
|
||||
auto gradW_user_mem = onednnUtils::loadDataToMklStream(
|
||||
*gradW, engine, stream, gradW_user_md, op_weights_bp_prim_desc.diff_weights_desc(), args[DNNL_ARG_DIFF_WEIGHTS]);
|
||||
|
||||
// gradB
|
||||
if (gradB != nullptr) {
|
||||
auto gradB_mkl_mem = dnnl::memory(gradB_mkl_md, engine, gradB->buffer());
|
||||
args[DNNL_ARG_DIFF_BIAS] = gradB_mkl_mem;
|
||||
}
|
||||
|
||||
// run backward data calculations
|
||||
dnnl::deconvolution_backward_data(op_data_bp_prim_desc).execute(stream, args);
|
||||
|
||||
if (gradOReorderW || gradOReorderD) args[DNNL_ARG_DIFF_DST] = gradO_mkl_memW;
|
||||
|
||||
// run backward weights calculations
|
||||
dnnl::deconvolution_backward_weights(op_weights_bp_prim_desc).execute(stream, args);
|
||||
|
||||
// reorder gradI if necessary
|
||||
if (op_data_bp_prim_desc.diff_src_desc() != gradI_user_mem.get_desc())
|
||||
dnnl::reorder(args[DNNL_ARG_DIFF_SRC], gradI_user_mem).execute(stream, args[DNNL_ARG_DIFF_SRC], gradI_user_mem);
|
||||
if (op_weights_bp_prim_desc.diff_weights_desc() != gradW_user_mem.get_desc())
|
||||
dnnl::reorder(args[DNNL_ARG_DIFF_WEIGHTS], gradW_user_mem)
|
||||
.execute(stream, args[DNNL_ARG_DIFF_WEIGHTS], gradW_user_mem);
|
||||
|
||||
stream.wait();
|
||||
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
PLATFORM_IMPL(deconv2d, ENGINE_CPU) {
|
||||
auto input = INPUT_VARIABLE(0); // [bS, iH, iW, iC] (NHWC) or [bS, iC, iH, iW] (NCHW)
|
||||
auto weights = INPUT_VARIABLE(1); // [kH, kW, oC, iC], [iC, oC, kH, kW], [iC, kH, kW, oC]
|
||||
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)
|
||||
|
||||
REQUIRE_TRUE(input->rankOf() == 4, 0,
|
||||
"CUSTOM DECONV2D_MKLDNN OP: rank of input array must be equal to 4, but got %i instead !",
|
||||
input->rankOf());
|
||||
REQUIRE_TRUE(weights->rankOf() == 4, 0,
|
||||
"CUSTOM DECONV2D_MKLDNN OP: rank of weights array must be equal to 4, but got %i instead !",
|
||||
weights->rankOf());
|
||||
|
||||
sd::LongType kH = INT_ARG(0) > 0 ? INT_ARG(0) : static_cast<sd::LongType>(weights->sizeAt(0)); // filter(kernel) height
|
||||
sd::LongType kW = INT_ARG(1) > 0 ? INT_ARG(1) : static_cast<sd::LongType>(weights->sizeAt(1)); // filter(kernel) width
|
||||
sd::LongType sH = INT_ARG(2); // strides height
|
||||
sd::LongType sW = INT_ARG(3); // strides width
|
||||
sd::LongType pH = INT_ARG(4); // paddings height
|
||||
sd::LongType pW = INT_ARG(5); // paddings width
|
||||
sd::LongType dH = INT_ARG(6); // dilations height
|
||||
sd::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, oC, iC], 1 - [iC, oC, kH, kW], 2 - [iC, kH, kW, oC]
|
||||
|
||||
sd::LongType bS, iC, iH, iW, oC, oH,
|
||||
oW; // batch size, input channels, input height/width, output channels, output height/width;
|
||||
sd::LongType indIOioC, indIiH, indWoC, indWiC, indWkH, indOoH; // corresponding indexes
|
||||
ConvolutionUtils::getSizesAndIndexesConv2d(isNCHW, wFormat, *input, *output, bS, iC, iH, iW, oC, oH, oW, indIOioC,
|
||||
indIiH, indWoC, indWiC, indWkH, indOoH);
|
||||
|
||||
std::vector<sd::LongType> expectedWeightsShape = ConvolutionUtils::expectWeightsShape(wFormat, kH, kW, oC, iC);
|
||||
REQUIRE_TRUE(weights->isSameShape(expectedWeightsShape), 0,
|
||||
"CUSTOM DECONV2D_MKLDNN 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 DECONV2D_MKLDNN OP: wrong shape of array with biases, expected rank, length: <=2, %i, but got "
|
||||
"%i, %i instead !",
|
||||
oC, bias->rankOf(), bias->lengthOf());
|
||||
|
||||
if (paddingMode) { // SAME
|
||||
// Note: we're intentionally swapping iH and oH, to calculated the padding for a"normal" conv (not deconv) forward
|
||||
// pass
|
||||
ConvolutionUtils::calcPadding2D(pH, pW, iH, iW, oH, oW, kH, kW, sH, sW, dH, dW);
|
||||
}
|
||||
|
||||
deconv2dMKLDNN(input, weights, bias, output, kH, kW, sH, sW, pH, pW, dH, dW, paddingMode, isNCHW, wFormat);
|
||||
|
||||
return sd::Status::OK;
|
||||
}
|
||||
|
||||
PLATFORM_CHECK(deconv2d, ENGINE_CPU) {
|
||||
auto input = INPUT_VARIABLE(0);
|
||||
auto weights = INPUT_VARIABLE(1);
|
||||
auto bias = block.width() > 2 ? INPUT_VARIABLE(2) : nullptr;
|
||||
|
||||
auto output = INPUT_VARIABLE(0);
|
||||
|
||||
int dH = INT_ARG(6); // dilations height
|
||||
int dW = INT_ARG(7); // dilations width
|
||||
int paddingMode = INT_ARG(8); // 0-VALID, 1-SAME
|
||||
|
||||
Requirements req("ONEDNN DECONV2d OP");
|
||||
req.expectTrue(block.isUseONEDNN(), IS_USE_ONEDNN_MSG) &&
|
||||
req.expectLessEq(makeInfoVariable(dH, "Dilation height"), 1) &&
|
||||
req.expectLessEq(makeInfoVariable(dW, "Dilation width"), 1) &&
|
||||
req.expectFalse(makeInfoVariable(paddingMode, "paddingMode")) &&
|
||||
req.expectTrue(makeInfoVariable(
|
||||
[input, weights, bias, output] {
|
||||
const DataType xType = input->dataType();
|
||||
const DataType wType = weights->dataType();
|
||||
const DataType zType = output->dataType();
|
||||
const DataType bType = bias != nullptr ? bias->dataType() : zType;
|
||||
return ((xType == DataType::FLOAT32 && wType == DataType::FLOAT32 &&
|
||||
bType == DataType::FLOAT32 && zType == DataType::FLOAT32) ||
|
||||
((xType == DataType::UINT8 || xType == DataType::INT8) && wType == DataType::INT8 &&
|
||||
(zType == DataType::UINT8 || zType == DataType::INT8 || zType == DataType::INT32 ||
|
||||
zType == DataType::FLOAT32) &&
|
||||
bType == zType));
|
||||
},
|
||||
TYPECHECK_MSG),
|
||||
NO_MSG);
|
||||
req.logTheSuccess();
|
||||
return req;
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
PLATFORM_IMPL(deconv2d_bp, ENGINE_CPU) {
|
||||
auto input = INPUT_VARIABLE(0); // [bS, iH, iW, iC] (NHWC) or [bS, iC, iH, iW] (NCDHW)
|
||||
auto weights = INPUT_VARIABLE(1); // [kH, kW, oC, iC], [iC, oC, kH, kW], [iC, kH, kW, oC]
|
||||
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] (NCDHW), epsilon_next
|
||||
|
||||
auto gradI = OUTPUT_VARIABLE(0); // [bS, iH, iW, iC] (NHWC) or [bS, iC, iH, iW] (NCDHW), gradI
|
||||
auto gradW = OUTPUT_VARIABLE(1); // [kH, kW, oC, iC], [iC, oC, kH, kW], [iC, kH, kW, oC]
|
||||
auto gradB = block.width() > 3 ? OUTPUT_VARIABLE(2) : nullptr; // [oC]
|
||||
|
||||
REQUIRE_TRUE(input->rankOf() == 4, 0,
|
||||
"CUSTOM DECONV2D_MKLDNN_BP OP: rank of input array must be equal to 4, but got %i instead !",
|
||||
input->rankOf());
|
||||
REQUIRE_TRUE(weights->rankOf() == 4, 0,
|
||||
"CUSTOM DECONV2D_MKLDNN_BP OP: rank of weights array must be equal to 4 , but got %i instead !",
|
||||
weights->rankOf());
|
||||
REQUIRE_TRUE(gradO->rankOf() == 4, 0,
|
||||
"CUSTOM DECONV2D_MKLDNN_BP OP: rank of output gradients (next epsilon) array must be equal to 4, but "
|
||||
"got %i instead !",
|
||||
gradO->rankOf());
|
||||
|
||||
sd::LongType kH = INT_ARG(0) > 0 ? INT_ARG(0) : static_cast<sd::LongType>(weights->sizeAt(0)); // filter(kernel) height
|
||||
sd::LongType kW = INT_ARG(1) > 0 ? INT_ARG(1) : static_cast<sd::LongType>(weights->sizeAt(1)); // filter(kernel) width
|
||||
sd::LongType sH = INT_ARG(2); // strides height
|
||||
sd::LongType sW = INT_ARG(3); // strides width
|
||||
sd::LongType pH = INT_ARG(4); // paddings height
|
||||
sd::LongType pW = INT_ARG(5); // paddings width
|
||||
sd::LongType dH = INT_ARG(6); // dilations height
|
||||
sd::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, oC, iC], 1 - [iC, oC, kH, kW], 2 - [iC, kH, kW, oC]
|
||||
|
||||
sd::LongType bS, iC, iH, iW, oC, oH,
|
||||
oW; // batch size, input channels, input height/width, output channels, output height/width;
|
||||
sd::LongType indIOioC, indIiH, indWoC, indWiC, indWkH, indOoH; // corresponding indexes
|
||||
ConvolutionUtils::getSizesAndIndexesConv2d(isNCHW, wFormat, *input, *gradO, bS, iC, iH, iW, oC, oH, oW, indIOioC,
|
||||
indIiH, indWoC, indWiC, indWkH, indOoH);
|
||||
|
||||
sd::LongType trueoH, trueoW; // true output height, width
|
||||
ConvolutionUtils::calcOutSizeDeconv2D(trueoH, trueoW, kH, kW, sH, sW, pH, pW, dH, dW, iH, iW, paddingMode);
|
||||
|
||||
std::vector<sd::LongType> expectedGradOShape =
|
||||
ShapeUtils::composeShapeUsingDimsAndIdx({bS, oC, trueoH, trueoW, 0, indIOioC, indOoH, indOoH + 1});
|
||||
std::vector<sd::LongType> expectedWeightsShape = ConvolutionUtils::expectWeightsShape(wFormat, kH, kW, oC, iC);
|
||||
REQUIRE_TRUE(gradO->isSameShape(expectedGradOShape), 0,
|
||||
"CUSTOM DECONV2D_MKLDNN_BP 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 DECONV2D_MKLDNN_BP 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 DECONV2D_MKLDNN_BP OP: wrong shape of array with biases, expected rank, length: <=2, %i, but "
|
||||
"got %i, %i instead !",
|
||||
oC, bias->rankOf(), bias->lengthOf());
|
||||
|
||||
if (paddingMode) { // SAME
|
||||
// Note: we're intentionally swapping iH and oH, to calculated the padding for a"normal" conv (not deconv) forward
|
||||
// pass
|
||||
ConvolutionUtils::calcPadding2D(pH, pW, iH, iW, oH, oW, kH, kW, sH, sW, dH, dW);
|
||||
}
|
||||
|
||||
deconv2dBpMKLDNN(input, weights, gradO, gradI, gradW, gradB, kH, kW, sH, sW, pH, pW, dH, dW, paddingMode, isNCHW,
|
||||
wFormat);
|
||||
|
||||
return sd::Status::OK;
|
||||
}
|
||||
|
||||
PLATFORM_CHECK(deconv2d_bp, ENGINE_CPU) {
|
||||
auto input = INPUT_VARIABLE(0); // [bS, iH, iW, iC] (NHWC) or [bS, iC, iH, iW] (NCDHW)
|
||||
auto weights = INPUT_VARIABLE(1); // [kH, kW, oC, iC], [iC, oC, kH, kW], [iC, kH, kW, oC]
|
||||
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] (NCDHW), epsilon_next
|
||||
|
||||
auto gradI = OUTPUT_VARIABLE(0); // [bS, iH, iW, iC] (NHWC) or [bS, iC, iH, iW] (NCDHW), gradI
|
||||
auto gradW = OUTPUT_VARIABLE(1); // [kH, kW, oC, iC], [iC, oC, kH, kW], [iC, kH, kW, oC]
|
||||
auto gradB = block.width() > 3 ? OUTPUT_VARIABLE(2) : nullptr; // [oC]
|
||||
|
||||
int dH = INT_ARG(6); // dilations height
|
||||
int dW = INT_ARG(7); // dilations width
|
||||
int paddingMode = INT_ARG(8); // 0-VALID, 1-SAME
|
||||
|
||||
Requirements req("ONEDNN DECONV2d_BP OP");
|
||||
req.expectTrue(block.isUseONEDNN(), IS_USE_ONEDNN_MSG) &&
|
||||
req.expectLessEq(makeInfoVariable(dH, "Dilation height"), 1) &&
|
||||
req.expectLessEq(makeInfoVariable(dW, "Dilation width"), 1) &&
|
||||
req.expectFalse(makeInfoVariable(paddingMode, "paddingMode")) &&
|
||||
req.expectTrue(makeInfoVariable(
|
||||
[input, weights, gradO, gradI, gradW, gradB] {
|
||||
const DataType xType = input->dataType();
|
||||
const DataType wType = weights->dataType();
|
||||
const DataType gradOType = gradO->dataType();
|
||||
|
||||
const DataType gradIType = gradI->dataType();
|
||||
const DataType gradWType = gradW->dataType();
|
||||
const DataType gradBType = gradB != nullptr ? gradB->dataType() : DataType::FLOAT32;
|
||||
return ((xType == DataType::FLOAT32 || xType == DataType::BFLOAT16) &&
|
||||
(wType == DataType::FLOAT32 || wType == DataType::BFLOAT16) &&
|
||||
(gradOType == DataType::FLOAT32 || gradOType == DataType::BFLOAT16) &&
|
||||
(gradIType == DataType::FLOAT32 || gradIType == DataType::BFLOAT16) &&
|
||||
(gradWType == DataType::FLOAT32 || gradWType == DataType::BFLOAT16) &&
|
||||
(gradBType == DataType::FLOAT32 || gradBType == DataType::BFLOAT16));
|
||||
},
|
||||
TYPECHECK_MSG),
|
||||
NO_MSG);
|
||||
req.logTheSuccess();
|
||||
return req;
|
||||
}
|
||||
|
||||
} // namespace platforms
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,232 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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 <helpers/MKLDNNStream.h>
|
||||
#include <ops/declarable/OpRegistrator.h>
|
||||
#include <ops/declarable/PlatformHelper.h>
|
||||
#include <ops/declarable/helpers/convolutions.h>
|
||||
#include <system/platform_boilerplate.h>
|
||||
|
||||
#include "mkldnnUtils.h"
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace platforms {
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
static void deconv2TFdBpMKLDNN(NDArray* weights, NDArray* gradO, NDArray* gradI, const sd::LongType bS, const sd::LongType iC,
|
||||
const sd::LongType iH, const sd::LongType iW, const sd::LongType oC, const sd::LongType oH, const sd::LongType oW, const sd::LongType kH,
|
||||
const sd::LongType kW, const sd::LongType sH, const sd::LongType sW, const sd::LongType pH, const sd::LongType pW, const sd::LongType dH,
|
||||
const sd::LongType dW, const bool isNCHW, const int wFormat) {
|
||||
// gradI [bS, iH, iW, iC], mkl doesn't support ndhwc format
|
||||
// weights [oC, iC, kH, kW] always, mkl doesn't support weights format [kH, kW, iC, oC]
|
||||
// gradO [bS, oH, oW, oC]
|
||||
|
||||
dnnl::memory::dims strides = {sH, sW};
|
||||
dnnl::memory::dims dilation = {dH - 1, dW - 1};
|
||||
dnnl::memory::dims padding = {pH, pW};
|
||||
dnnl::memory::dims padding_r = {(oH - 1) * sH - iH + kH - pH, (oW - 1) * sW - iW + kW - pW};
|
||||
|
||||
// weights type
|
||||
dnnl::memory::data_type wType =
|
||||
weights->dataType() == DataType::FLOAT32 ? dnnl::memory::data_type::f32 : dnnl::memory::data_type::bf16;
|
||||
// gradO type
|
||||
dnnl::memory::data_type gradOType =
|
||||
gradO->dataType() == DataType::FLOAT32 ? dnnl::memory::data_type::f32 : dnnl::memory::data_type::bf16;
|
||||
// gradI type
|
||||
dnnl::memory::data_type gradIType =
|
||||
gradI->dataType() == DataType::FLOAT32 ? dnnl::memory::data_type::f32 : dnnl::memory::data_type::bf16;
|
||||
|
||||
dnnl::memory::format_tag xFormatMkl = isNCHW ? dnnl::memory::format_tag::nchw : dnnl::memory::format_tag::nhwc;
|
||||
dnnl::memory::format_tag wFormatMkl = dnnl::memory::format_tag::oihw;
|
||||
|
||||
dnnl::memory::dims xDims = {bS, iC, iH, iW};
|
||||
dnnl::memory::dims wDims = {oC, iC, kH, kW};
|
||||
dnnl::memory::dims zDims = {bS, oC, oH, oW};
|
||||
|
||||
// memory descriptors for arrays
|
||||
|
||||
// input
|
||||
dnnl::memory::desc x_mkl_md = dnnl::memory::desc(xDims, gradOType, dnnl::memory::format_tag::any);
|
||||
|
||||
// weights
|
||||
dnnl::memory::desc w_mkl_md = dnnl::memory::desc(wDims, wType, dnnl::memory::format_tag::any);
|
||||
dnnl::memory::desc w_user_md = dnnl::memory::desc(wDims, wType, wFormatMkl);
|
||||
onednnUtils::setBlockStrides(*weights, w_user_md, {3, 2, 0, 1}); // permute [kH, kW, iC, oC] -> [oC, iC, kH, kW]
|
||||
|
||||
// gradO
|
||||
dnnl::memory::desc gradO_mkl_md = dnnl::memory::desc(zDims, gradOType, dnnl::memory::format_tag::any);
|
||||
dnnl::memory::desc gradO_user_md = dnnl::memory::desc(zDims, gradOType, xFormatMkl);
|
||||
onednnUtils::setBlockStrides(*gradO, gradO_user_md);
|
||||
|
||||
// gradI
|
||||
dnnl::memory::desc gradI_mkl_md = dnnl::memory::desc(xDims, gradIType, dnnl::memory::format_tag::any);
|
||||
dnnl::memory::desc gradI_user_md = dnnl::memory::desc(xDims, gradIType, xFormatMkl);
|
||||
onednnUtils::setBlockStrides(*gradI, gradI_user_md);
|
||||
|
||||
auto engine = onednnUtils::getEngine(LaunchContext::defaultContext()->engine());
|
||||
|
||||
// forward primitive description
|
||||
dnnl::convolution_forward::desc op_ff_desc(dnnl::prop_kind::forward_inference, dnnl::algorithm::convolution_auto,
|
||||
x_mkl_md, w_mkl_md, gradO_mkl_md, strides, dilation, padding, padding_r);
|
||||
dnnl::convolution_forward::primitive_desc op_ff_prim_desc(op_ff_desc, engine);
|
||||
|
||||
// backward data primitive description
|
||||
dnnl::convolution_backward_data::desc op_data_bp_desc(dnnl::algorithm::convolution_auto, gradI_mkl_md, w_mkl_md,
|
||||
gradO_mkl_md, strides, dilation, padding, padding_r);
|
||||
dnnl::convolution_backward_data::primitive_desc op_data_bp_prim_desc(op_data_bp_desc, engine, op_ff_prim_desc);
|
||||
|
||||
// arguments (memory buffers) necessary for calculations
|
||||
std::unordered_map<int, dnnl::memory> args;
|
||||
|
||||
dnnl::stream stream(engine);
|
||||
|
||||
// provide memory buffers and check whether reorder is required
|
||||
|
||||
// weights
|
||||
onednnUtils::loadDataToMklStream(*weights, engine, stream, w_user_md, op_data_bp_prim_desc.weights_desc(),
|
||||
args[DNNL_ARG_WEIGHTS]);
|
||||
|
||||
// gradO
|
||||
onednnUtils::loadDataToMklStream(*gradO, engine, stream, gradO_user_md, op_data_bp_prim_desc.diff_dst_desc(),
|
||||
args[DNNL_ARG_DIFF_DST]);
|
||||
|
||||
// gradI
|
||||
auto gradI_user_mem = onednnUtils::loadDataToMklStream(*gradI, engine, stream, gradI_user_md,
|
||||
op_data_bp_prim_desc.diff_src_desc(), args[DNNL_ARG_DIFF_SRC]);
|
||||
|
||||
// run backward data calculations
|
||||
dnnl::convolution_backward_data(op_data_bp_prim_desc).execute(stream, args);
|
||||
|
||||
// reorder gradI if necessary
|
||||
if (op_data_bp_prim_desc.diff_src_desc() != gradI_user_mem.get_desc())
|
||||
dnnl::reorder(args[DNNL_ARG_DIFF_SRC], gradI_user_mem).execute(stream, args[DNNL_ARG_DIFF_SRC], gradI_user_mem);
|
||||
|
||||
stream.wait();
|
||||
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
PLATFORM_IMPL(deconv2d_tf, ENGINE_CPU) {
|
||||
auto gradO = INPUT_VARIABLE(2); // [bS, oH, oW, oC] (NHWC) or [bS, oC, oH, oW] (NCHW), epsilon_next
|
||||
auto weights = INPUT_VARIABLE(1); // [kH, kW, iC, oC], [oC, iC, kH, kW], [oC, kH, kW, iC]
|
||||
auto gradIShape = INPUT_VARIABLE(0); // [4] - shape of input of conv2d (that is shape of gradI)
|
||||
|
||||
auto gradI = OUTPUT_VARIABLE(0); // [bS, iH, iW, iC] (NHWC) or [bS, iC, iH, iW] (NCHW), epsilon
|
||||
|
||||
sd::LongType kH = INT_ARG(0) > 0 ? INT_ARG(0) : static_cast<sd::LongType>(weights->sizeAt(0)); // filter(kernel) height
|
||||
sd::LongType kW = INT_ARG(1) > 0 ? INT_ARG(1) : static_cast<sd::LongType>(weights->sizeAt(1)); // filter(kernel) width
|
||||
sd::LongType sH = INT_ARG(2); // strides height
|
||||
sd::LongType sW = INT_ARG(3); // strides width
|
||||
sd::LongType pH = INT_ARG(4); // paddings height
|
||||
sd::LongType pW = INT_ARG(5); // paddings width
|
||||
sd::LongType dH = INT_ARG(6); // dilations height
|
||||
sd::LongType dW = INT_ARG(7); // dilations width
|
||||
int isSameMode = 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, oC], 1 - [oC, iC, kH, kW], 2 - [oC, kH, kW, iC]
|
||||
|
||||
const sd::LongType rank = gradO->rankOf();
|
||||
|
||||
REQUIRE_TRUE(weights->rankOf() == rank, 0,
|
||||
"CUSTOM DECONV2D_TF MKLDNN OP: rank of weights array must be equal to 4, but got %i instead !",
|
||||
weights->rankOf());
|
||||
REQUIRE_TRUE(gradIShape->rankOf() == 1, 0,
|
||||
"CUSTOM DECONV2D_TF MKLDNN OP: rank of array with output shape must be equal to 1, but got %i instead !",
|
||||
gradIShape->rankOf());
|
||||
REQUIRE_TRUE(
|
||||
gradIShape->lengthOf() == rank, 0,
|
||||
"CUSTOM DECONV2D_TF MKLDNN OP: length of array with output shape must be equal to 4, but got %i instead !",
|
||||
gradIShape->lengthOf());
|
||||
|
||||
int indIOioC, indIiH, indWoC(3), indOoH;
|
||||
if (!isNCHW) {
|
||||
indIOioC = 3;
|
||||
indIiH = 1;
|
||||
indOoH = 1;
|
||||
} else {
|
||||
indIOioC = 1;
|
||||
indIiH = 2;
|
||||
indOoH = 2;
|
||||
}
|
||||
|
||||
std::vector<sd::LongType> gradIShapeVector = gradIShape->template asVectorT<sd::LongType>();
|
||||
|
||||
const sd::LongType bS = gradIShapeVector[0]; // batch size
|
||||
const sd::LongType iH = gradIShapeVector[indIiH]; // input height
|
||||
const sd::LongType iW = gradIShapeVector[indIiH + 1]; // input width
|
||||
const sd::LongType iC = gradIShapeVector[indIOioC]; // input channels
|
||||
const sd::LongType oC = weights->sizeAt(indWoC); // output channels
|
||||
const sd::LongType oH = gradO->sizeAt(indOoH); // input height
|
||||
const sd::LongType oW = gradO->sizeAt(indOoH); // input width
|
||||
|
||||
sd::LongType trueoH, trueoW; // true output height, width
|
||||
ConvolutionUtils::calcOutSizePool2D(trueoH, trueoW, kH, kW, sH, sW, pH, pW, dH, dW, iH, iW, isSameMode);
|
||||
|
||||
std::vector<sd::LongType> expectedGradOShape =
|
||||
ShapeUtils::composeShapeUsingDimsAndIdx({bS, oC, trueoH, trueoW, 0, indIOioC, indOoH, indOoH + 1});
|
||||
std::vector<sd::LongType> expectedWeightsShape = {kH, kW, iC, oC};
|
||||
REQUIRE_TRUE(gradO->isSameShape(expectedGradOShape), 0,
|
||||
"CUSTOM DECONV2D_TF MKLDNN OP: wrong shape of input array, basing on array with output shape expected "
|
||||
"is %s, but got %s instead !",
|
||||
ShapeUtils::shapeAsString(expectedGradOShape).c_str(), ShapeUtils::shapeAsString(gradO).c_str());
|
||||
REQUIRE_TRUE(weights->isSameShape(expectedWeightsShape), 0,
|
||||
"CUSTOM DECONV2D_TF MKLDNN OP: wrong shape of weights array, expected is %s, but got %s instead !",
|
||||
ShapeUtils::shapeAsString(expectedWeightsShape).c_str(), ShapeUtils::shapeAsString(weights).c_str());
|
||||
|
||||
if (isSameMode) // SAME
|
||||
ConvolutionUtils::calcPadding2D(pH, pW, oH, oW, iH, iW, kH, kW, sH, sW, dH, dW);
|
||||
|
||||
// // mkl supports only [oC, iC, kH, kW] for weights
|
||||
|
||||
|
||||
deconv2TFdBpMKLDNN(weights, gradO, gradI, bS, iC, iH, iW, oC, oH, oW, kH, kW, sH, sW, pH, pW, dH, dW, isNCHW,
|
||||
wFormat);
|
||||
|
||||
return sd::Status::OK;
|
||||
}
|
||||
|
||||
PLATFORM_CHECK(deconv2d_tf, ENGINE_CPU) {
|
||||
auto weights = INPUT_VARIABLE(1); // [kH, kW, iC, oC] always
|
||||
auto gradO = INPUT_VARIABLE(2); // [bS, oH, oW, oC] (NHWC) or [bS, oC, oH, oW] (NCDHW), epsilon_next
|
||||
auto gradI = OUTPUT_VARIABLE(0); // [bS, iH, iW, iC] (NHWC) or [bS, iC, iH, iW] (NCDHW), gradI
|
||||
Requirements req("ONEDNN DECONV2d_TF OP");
|
||||
req.expectTrue(block.isUseONEDNN(), IS_USE_ONEDNN_MSG) &&
|
||||
req.expectTrue(makeInfoVariable(
|
||||
[weights, gradI, gradO] {
|
||||
const DataType wType = weights->dataType();
|
||||
const DataType gradOType = gradO->dataType();
|
||||
const DataType gradIType = gradI->dataType();
|
||||
return ((wType == DataType::FLOAT32 || wType == DataType::BFLOAT16) &&
|
||||
(gradOType == DataType::FLOAT32 || gradOType == DataType::BFLOAT16) &&
|
||||
(gradIType == DataType::FLOAT32 || gradIType == DataType::BFLOAT16));
|
||||
},
|
||||
TYPECHECK_MSG),
|
||||
NO_MSG);
|
||||
req.logTheSuccess();
|
||||
return req;
|
||||
}
|
||||
|
||||
} // namespace platforms
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,546 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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 <helpers/MKLDNNStream.h>
|
||||
#include <ops/declarable/OpRegistrator.h>
|
||||
#include <ops/declarable/PlatformHelper.h>
|
||||
#include <ops/declarable/helpers/convolutions.h>
|
||||
#include <system/platform_boilerplate.h>
|
||||
|
||||
#include "mkldnnUtils.h"
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace platforms {
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
static void deconv3dMKLDNN(NDArray* input, NDArray* weights, NDArray* bias, NDArray* output,
|
||||
const sd::LongType kD, const sd::LongType kH, const sd::LongType kW, const sd::LongType sD, const sd::LongType sH, const sd::LongType sW,
|
||||
const sd::LongType pD, const sd::LongType pH, const sd::LongType pW, const sd::LongType dD, const sd::LongType dH, const sd::LongType dW,
|
||||
const bool isNCDHW, const int wFormat) {
|
||||
// mkl supports weights in [oC, iC, kD, kH, kW] only
|
||||
|
||||
sd::LongType bS, iC, iD, iH, iW, oC, oD, oH,
|
||||
oW; // batch size, input channels, input depth/height/width, output channels, output depth/height/width;
|
||||
sd::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, indWoC, indWiC, indWkD);
|
||||
|
||||
dnnl::memory::dims strides = {sD, sH, sW};
|
||||
dnnl::memory::dims padding = {pD, pH, pW};
|
||||
dnnl::memory::dims padding_r = {(iD - 1) * sD - oD + kD - pD, (iH - 1) * sH - oH + kH - pH,
|
||||
(iW - 1) * sW - oW + kW - pW};
|
||||
dnnl::memory::dims dilation = {dD - 1, dH - 1, dW - 1};
|
||||
|
||||
std::vector<int> permut;
|
||||
if (0 == wFormat)
|
||||
permut = {3, 4, 0, 1, 2}; // [kD, kH, kW, oC, iC] -> [oC, iC, kD, kH, kW]
|
||||
else if (1 == wFormat)
|
||||
permut = {1, 0, 2, 3, 4}; // [iC, oC, kD, kH, kW] -> [oC, iC, kD, kH, kW]
|
||||
else
|
||||
permut = {4, 0, 1, 2, 3}; // [iC, kD, kH, kW, oC] -> [oC, iC, kD, kH, kW]
|
||||
|
||||
// input type
|
||||
dnnl::memory::data_type xType;
|
||||
if (input->dataType() == DataType::FLOAT32)
|
||||
xType = dnnl::memory::data_type::f32;
|
||||
else if (input->dataType() == DataType::HALF)
|
||||
xType = dnnl::memory::data_type::f16;
|
||||
else if (input->dataType() == DataType::UINT8)
|
||||
xType = dnnl::memory::data_type::u8;
|
||||
else
|
||||
xType = dnnl::memory::data_type::s8;
|
||||
|
||||
// weights type
|
||||
dnnl::memory::data_type wType = xType;
|
||||
if (xType == dnnl::memory::data_type::u8) wType = dnnl::memory::data_type::s8;
|
||||
|
||||
// output and bias type (have the same types)
|
||||
dnnl::memory::data_type zType;
|
||||
if (output->dataType() == DataType::FLOAT32)
|
||||
zType = dnnl::memory::data_type::f32;
|
||||
else if (output->dataType() == DataType::HALF)
|
||||
zType = dnnl::memory::data_type::f16;
|
||||
else if (output->dataType() == DataType::UINT8)
|
||||
zType = dnnl::memory::data_type::u8;
|
||||
else if (output->dataType() == DataType::INT8)
|
||||
zType = dnnl::memory::data_type::s8;
|
||||
else
|
||||
zType = dnnl::memory::data_type::s32;
|
||||
|
||||
dnnl::memory::format_tag xFormatMkl = isNCDHW ? dnnl::memory::format_tag::ncdhw : dnnl::memory::format_tag::ndhwc;
|
||||
dnnl::memory::format_tag wFormatMkl = dnnl::memory::format_tag::oidhw;
|
||||
|
||||
dnnl::memory::dims xDims = {bS, iC, iD, iH, iW};
|
||||
dnnl::memory::dims wDims = {oC, iC, kD, kH, kW};
|
||||
dnnl::memory::dims zDims = {bS, oC, oD, oH, oW};
|
||||
|
||||
// memory descriptors for arrays
|
||||
|
||||
// input
|
||||
dnnl::memory::desc x_mkl_md = dnnl::memory::desc(xDims, xType, dnnl::memory::format_tag::any);
|
||||
dnnl::memory::desc x_user_md = dnnl::memory::desc(xDims, xType, xFormatMkl);
|
||||
onednnUtils::setBlockStrides(*input, x_user_md);
|
||||
|
||||
// weights
|
||||
dnnl::memory::desc w_mkl_md = dnnl::memory::desc(wDims, wType, dnnl::memory::format_tag::any);
|
||||
dnnl::memory::desc w_user_md = dnnl::memory::desc(wDims, wType, wFormatMkl);
|
||||
onednnUtils::setBlockStrides(*weights, w_user_md, permut);
|
||||
|
||||
// bias
|
||||
dnnl::memory::desc b_mkl_md;
|
||||
if (bias != nullptr) b_mkl_md = dnnl::memory::desc({oC}, zType, dnnl::memory::format_tag::x);
|
||||
|
||||
// output
|
||||
dnnl::memory::desc z_mkl_md = dnnl::memory::desc(zDims, zType, dnnl::memory::format_tag::any);
|
||||
dnnl::memory::desc z_user_md = dnnl::memory::desc(zDims, zType, xFormatMkl);
|
||||
onednnUtils::setBlockStrides(*output, z_user_md);
|
||||
|
||||
auto engine = onednnUtils::getEngine(LaunchContext::defaultContext()->engine());
|
||||
|
||||
// operation primitive description
|
||||
dnnl::deconvolution_forward::desc op_desc(dnnl::prop_kind::forward_inference, dnnl::algorithm::deconvolution_direct,
|
||||
x_mkl_md, w_mkl_md, b_mkl_md, z_mkl_md, strides, dilation, padding,
|
||||
padding_r);
|
||||
dnnl::deconvolution_forward::primitive_desc op_prim_desc(op_desc, engine);
|
||||
|
||||
// arguments (memory buffers) necessary for calculations
|
||||
std::unordered_map<int, dnnl::memory> args;
|
||||
|
||||
dnnl::stream stream(engine);
|
||||
|
||||
// provide memory buffers and check whether reorder is required
|
||||
|
||||
// input
|
||||
onednnUtils::loadDataToMklStream(*input, engine, stream, x_user_md, op_prim_desc.src_desc(), args[DNNL_ARG_SRC]);
|
||||
|
||||
// weights
|
||||
onednnUtils::loadDataToMklStream(*weights, engine, stream, w_user_md, op_prim_desc.weights_desc(),
|
||||
args[DNNL_ARG_WEIGHTS]);
|
||||
|
||||
// bias
|
||||
if (bias != nullptr) {
|
||||
auto b_mkl_mem = dnnl::memory(b_mkl_md, engine, const_cast<void*>(bias->buffer()));
|
||||
args[DNNL_ARG_BIAS] = b_mkl_mem;
|
||||
}
|
||||
|
||||
// output
|
||||
auto z_user_mem =
|
||||
onednnUtils::loadDataToMklStream(*output, engine, stream, z_user_md, op_prim_desc.dst_desc(), args[DNNL_ARG_DST]);
|
||||
|
||||
// run calculations
|
||||
dnnl::deconvolution_forward(op_prim_desc).execute(stream, args);
|
||||
|
||||
// reorder outputs if necessary
|
||||
if (op_prim_desc.dst_desc() != z_user_mem.get_desc())
|
||||
dnnl::reorder(args[DNNL_ARG_DST], z_user_mem).execute(stream, args[DNNL_ARG_DST], z_user_mem);
|
||||
|
||||
stream.wait();
|
||||
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
static void deconv3dBackPropMKLDNN(NDArray* input, NDArray* weights, NDArray* gradO, NDArray* gradI,
|
||||
NDArray* gradW, NDArray* gradB, const sd::LongType kD, const sd::LongType kH, const sd::LongType kW,
|
||||
const sd::LongType sD, const sd::LongType sH, const sd::LongType sW, const sd::LongType pD, const sd::LongType pH, const sd::LongType pW,
|
||||
const sd::LongType dD, const sd::LongType dH, const int dW, const bool isNCDHW, const int wFormat) {
|
||||
// mkl supports weights/gradW in [oC, iC, kD, kH, kW] format only
|
||||
|
||||
sd::LongType bS, iC, iD, iH, iW, oC, oD, oH,
|
||||
oW; // batch size, input channels, input depth/height/width, output channels, output depth/height/width;
|
||||
sd::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, indWoC, indWiC, indWkD);
|
||||
|
||||
dnnl::memory::dims strides = {sD, sH, sW};
|
||||
dnnl::memory::dims padding = {pD, pH, pW};
|
||||
dnnl::memory::dims padding_r = {(iD - 1) * sD - oD + kD - pD, (iH - 1) * sH - oH + kH - pH,
|
||||
(iW - 1) * sW - oW + kW - pW};
|
||||
dnnl::memory::dims dilation = {dD - 1, dH - 1, dW - 1};
|
||||
|
||||
std::vector<int> permut;
|
||||
if (0 == wFormat)
|
||||
permut = {3, 4, 0, 1, 2}; // [kD, kH, kW, oC, iC] -> [oC, iC, kD, kH, kW]
|
||||
else if (1 == wFormat)
|
||||
permut = {1, 0, 2, 3, 4}; // [iC, oC, kD, kH, kW] -> [oC, iC, kD, kH, kW]
|
||||
else
|
||||
permut = {4, 0, 1, 2, 3}; // [iC, kD, kH, kW, oC] -> [oC, iC, kD, kH, kW]
|
||||
|
||||
// input type
|
||||
dnnl::memory::data_type xType =
|
||||
input->dataType() == DataType::FLOAT32 ? dnnl::memory::data_type::f32 : dnnl::memory::data_type::bf16;
|
||||
// weights type
|
||||
dnnl::memory::data_type wType =
|
||||
weights->dataType() == DataType::FLOAT32 ? dnnl::memory::data_type::f32 : dnnl::memory::data_type::bf16;
|
||||
// gradO type
|
||||
dnnl::memory::data_type gradOType =
|
||||
gradO->dataType() == DataType::FLOAT32 ? dnnl::memory::data_type::f32 : dnnl::memory::data_type::bf16;
|
||||
// gradI type
|
||||
dnnl::memory::data_type gradIType =
|
||||
gradI->dataType() == DataType::FLOAT32 ? dnnl::memory::data_type::f32 : dnnl::memory::data_type::bf16;
|
||||
// gradW type
|
||||
dnnl::memory::data_type gradWType =
|
||||
gradW->dataType() == DataType::FLOAT32 ? dnnl::memory::data_type::f32 : dnnl::memory::data_type::bf16;
|
||||
// gradB type
|
||||
dnnl::memory::data_type gradBType =
|
||||
gradB != nullptr
|
||||
? (gradB->dataType() == DataType::FLOAT32 ? dnnl::memory::data_type::f32 : dnnl::memory::data_type::bf16)
|
||||
: dnnl::memory::data_type::f32;
|
||||
|
||||
dnnl::memory::format_tag xFormatMkl = isNCDHW ? dnnl::memory::format_tag::ncdhw : dnnl::memory::format_tag::ndhwc;
|
||||
dnnl::memory::format_tag wFormatMkl = dnnl::memory::format_tag::oidhw;
|
||||
|
||||
dnnl::memory::dims xDims = {bS, iC, iD, iH, iW};
|
||||
dnnl::memory::dims wDims = {oC, iC, kD, kH, kW};
|
||||
dnnl::memory::dims zDims = {bS, oC, oD, oH, oW};
|
||||
|
||||
// memory descriptors for arrays
|
||||
|
||||
// input
|
||||
dnnl::memory::desc x_mkl_md = dnnl::memory::desc(xDims, xType, dnnl::memory::format_tag::any);
|
||||
dnnl::memory::desc x_user_md = dnnl::memory::desc(xDims, xType, xFormatMkl);
|
||||
onednnUtils::setBlockStrides(*input, x_user_md);
|
||||
|
||||
// weights
|
||||
dnnl::memory::desc w_mkl_md = dnnl::memory::desc(wDims, wType, dnnl::memory::format_tag::any);
|
||||
dnnl::memory::desc w_user_md = dnnl::memory::desc(wDims, wType, wFormatMkl);
|
||||
onednnUtils::setBlockStrides(*weights, w_user_md, permut);
|
||||
|
||||
// gradO
|
||||
dnnl::memory::desc gradO_mkl_md = dnnl::memory::desc(zDims, gradOType, dnnl::memory::format_tag::any);
|
||||
dnnl::memory::desc gradO_user_md = dnnl::memory::desc(zDims, gradOType, xFormatMkl);
|
||||
onednnUtils::setBlockStrides(*gradO, gradO_user_md);
|
||||
|
||||
// gradI
|
||||
dnnl::memory::desc gradI_mkl_md = dnnl::memory::desc(xDims, gradIType, dnnl::memory::format_tag::any);
|
||||
dnnl::memory::desc gradI_user_md = dnnl::memory::desc(xDims, gradIType, xFormatMkl);
|
||||
onednnUtils::setBlockStrides(*gradI, gradI_user_md);
|
||||
|
||||
// gradW
|
||||
dnnl::memory::desc gradW_mkl_md = dnnl::memory::desc(wDims, gradWType, dnnl::memory::format_tag::any);
|
||||
dnnl::memory::desc gradW_user_md = dnnl::memory::desc(wDims, gradWType, wFormatMkl);
|
||||
onednnUtils::setBlockStrides(*gradW, gradW_user_md, permut);
|
||||
|
||||
// gradB
|
||||
dnnl::memory::desc gradB_mkl_md;
|
||||
if (gradB != nullptr) gradB_mkl_md = dnnl::memory::desc({oC}, gradBType, dnnl::memory::format_tag::x);
|
||||
|
||||
auto engine = onednnUtils::getEngine(LaunchContext::defaultContext()->engine());
|
||||
|
||||
// forward primitive description
|
||||
dnnl::deconvolution_forward::desc op_ff_desc(dnnl::prop_kind::forward_inference,
|
||||
dnnl::algorithm::deconvolution_direct, x_mkl_md, w_mkl_md, gradB_mkl_md,
|
||||
gradO_mkl_md, strides, dilation, padding, padding_r);
|
||||
dnnl::deconvolution_forward::primitive_desc op_ff_prim_desc(op_ff_desc, engine);
|
||||
|
||||
// backward data primitive description
|
||||
dnnl::deconvolution_backward_data::desc op_data_bp_desc(dnnl::algorithm::deconvolution_direct, gradI_mkl_md, w_mkl_md,
|
||||
gradO_mkl_md, strides, dilation, padding, padding_r);
|
||||
dnnl::deconvolution_backward_data::primitive_desc op_data_bp_prim_desc(op_data_bp_desc, engine, op_ff_prim_desc);
|
||||
|
||||
// backward weights primitive description
|
||||
dnnl::deconvolution_backward_weights::desc op_weights_bp_desc(dnnl::algorithm::deconvolution_direct, x_mkl_md,
|
||||
gradW_mkl_md, gradB_mkl_md, gradO_mkl_md, strides,
|
||||
dilation, padding, padding_r);
|
||||
dnnl::deconvolution_backward_weights::primitive_desc op_weights_bp_prim_desc(op_weights_bp_desc, engine,
|
||||
op_ff_prim_desc);
|
||||
|
||||
// arguments (memory buffers) necessary for calculations
|
||||
std::unordered_map<int, dnnl::memory> args;
|
||||
|
||||
dnnl::stream stream(engine);
|
||||
|
||||
// provide memory buffers and check whether reorder is required
|
||||
|
||||
// input
|
||||
onednnUtils::loadDataToMklStream(*input, engine, stream, x_user_md, op_weights_bp_prim_desc.src_desc(),
|
||||
args[DNNL_ARG_SRC]);
|
||||
|
||||
// weights
|
||||
onednnUtils::loadDataToMklStream(*weights, engine, stream, w_user_md, op_data_bp_prim_desc.weights_desc(),
|
||||
args[DNNL_ARG_WEIGHTS]);
|
||||
|
||||
// gradO
|
||||
auto gradO_user_mem = dnnl::memory(gradO_user_md, engine, const_cast<void*>(gradO->buffer()));
|
||||
const bool gradOReorderW = op_weights_bp_prim_desc.diff_dst_desc() != gradO_user_mem.get_desc();
|
||||
const bool gradOReorderD = op_data_bp_prim_desc.diff_dst_desc() != gradO_user_mem.get_desc();
|
||||
auto gradO_mkl_memW = gradOReorderW ? dnnl::memory(op_weights_bp_prim_desc.diff_dst_desc(), engine) : gradO_user_mem;
|
||||
auto gradO_mkl_memD = gradOReorderD ? dnnl::memory(op_data_bp_prim_desc.diff_dst_desc(), engine) : gradO_user_mem;
|
||||
if (gradOReorderW) dnnl::reorder(gradO_user_mem, gradO_mkl_memW).execute(stream, gradO_user_mem, gradO_mkl_memW);
|
||||
if (gradOReorderD) dnnl::reorder(gradO_user_mem, gradO_mkl_memD).execute(stream, gradO_user_mem, gradO_mkl_memD);
|
||||
args[DNNL_ARG_DIFF_DST] = gradO_mkl_memD;
|
||||
|
||||
// gradI
|
||||
auto gradI_user_mem = onednnUtils::loadDataToMklStream(*gradI, engine, stream, gradI_user_md,
|
||||
op_data_bp_prim_desc.diff_src_desc(), args[DNNL_ARG_DIFF_SRC]);
|
||||
|
||||
// gradW
|
||||
auto gradW_user_mem = onednnUtils::loadDataToMklStream(
|
||||
*gradW, engine, stream, gradW_user_md, op_weights_bp_prim_desc.diff_weights_desc(), args[DNNL_ARG_DIFF_WEIGHTS]);
|
||||
|
||||
// gradB
|
||||
if (gradB != nullptr) {
|
||||
auto gradB_mkl_mem = dnnl::memory(gradB_mkl_md, engine, gradB->buffer());
|
||||
args[DNNL_ARG_DIFF_BIAS] = gradB_mkl_mem;
|
||||
}
|
||||
|
||||
// run backward data calculations
|
||||
dnnl::deconvolution_backward_data(op_data_bp_prim_desc).execute(stream, args);
|
||||
|
||||
if (gradOReorderW || gradOReorderD) args[DNNL_ARG_DIFF_DST] = gradO_mkl_memW;
|
||||
|
||||
// run backward weights calculations
|
||||
dnnl::deconvolution_backward_weights(op_weights_bp_prim_desc).execute(stream, args);
|
||||
|
||||
// reorder gradI if necessary
|
||||
if (op_data_bp_prim_desc.diff_src_desc() != gradI_user_mem.get_desc())
|
||||
dnnl::reorder(args[DNNL_ARG_DIFF_SRC], gradI_user_mem).execute(stream, args[DNNL_ARG_DIFF_SRC], gradI_user_mem);
|
||||
if (op_weights_bp_prim_desc.diff_weights_desc() != gradW_user_mem.get_desc())
|
||||
dnnl::reorder(args[DNNL_ARG_DIFF_WEIGHTS], gradW_user_mem)
|
||||
.execute(stream, args[DNNL_ARG_DIFF_WEIGHTS], gradW_user_mem);
|
||||
|
||||
stream.wait();
|
||||
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
PLATFORM_IMPL(deconv3d, ENGINE_CPU) {
|
||||
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, oC, iC], [iC, oC, kD, kH, kW], [iC, kD, kH, kW, oC]
|
||||
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,
|
||||
"CUSTOM DECONV3D_MKLDNN OP: rank of input array must be equal to 5, but got %i instead !",
|
||||
input->rankOf());
|
||||
REQUIRE_TRUE(weights->rankOf() == 5, 0,
|
||||
"CUSTOM DECONV3D_MKLDNN OP: rank of weights array must be equal to 5, but got %i instead !",
|
||||
weights->rankOf());
|
||||
|
||||
sd::LongType kD = INT_ARG(0) > 0 ? INT_ARG(0) : static_cast<sd::LongType>(weights->sizeAt(0)); // filter(kernel) depth
|
||||
sd::LongType kH = INT_ARG(1) > 0 ? INT_ARG(1) : static_cast<sd::LongType>(weights->sizeAt(1)); // filter(kernel) height
|
||||
sd::LongType kW = INT_ARG(2) > 0 ? INT_ARG(2) : static_cast<sd::LongType>(weights->sizeAt(2)); // filter(kernel) width
|
||||
sd::LongType sD = INT_ARG(3); // strides depth
|
||||
sd::LongType sH = INT_ARG(4); // strides height
|
||||
sd::LongType sW = INT_ARG(5); // strides width
|
||||
sd::LongType pD = INT_ARG(6); // paddings depth
|
||||
sd::LongType pH = INT_ARG(7); // paddings height
|
||||
sd::LongType pW = INT_ARG(8); // paddings width
|
||||
sd::LongType dD = INT_ARG(9); // dilations depth
|
||||
sd::LongType dH = INT_ARG(10); // dilations height
|
||||
sd::LongType dW = INT_ARG(11); // dilations width
|
||||
int isSameMode = 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, oC, iC], 1 - [iC, oC, kD, kH, kW], 2 - [iC, kD, kH, kW, oC]
|
||||
|
||||
sd::LongType bS, iC, iD, iH, iW, oC, oD, oH,
|
||||
oW; // batch size, input channels, input depth/height/width, output channels, output depth/height/width;
|
||||
sd::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, indWoC, indWiC, indWkD);
|
||||
|
||||
std::vector<sd::LongType> expectedWeightsShape = ConvolutionUtils::expectWeightsShape(wFormat, kD, kH, kW, oC, iC);
|
||||
REQUIRE_TRUE(weights->isSameShape(expectedWeightsShape), 0,
|
||||
"CUSTOM DECONV3D_MKLDNN 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 DECONV3D_MKLDNN OP: wrong shape of array with biases, expected rank, length: <=2, %i, but got "
|
||||
"%i, %i instead !",
|
||||
oC, bias->rankOf(), bias->lengthOf());
|
||||
|
||||
if (isSameMode) { // SAME
|
||||
// Note: we're intentionally swapping iH and oH, to calculated the padding for a"normal" conv (not deconv) forward
|
||||
// pass
|
||||
ConvolutionUtils::calcPadding3D(pD, pH, pW, iD, iH, iW, oD, oH, oW, kD, kH, kW, sD, sH, sW, dD, dH, dW);
|
||||
}
|
||||
|
||||
deconv3dMKLDNN(input, weights, bias, output, kD, kH, kW, sD, sH, sW, pD, pH, pW, dD, dH, dW, isNCDHW, wFormat);
|
||||
|
||||
return sd::Status::OK;
|
||||
}
|
||||
|
||||
PLATFORM_CHECK(deconv3d, ENGINE_CPU) {
|
||||
auto input = INPUT_VARIABLE(0);
|
||||
auto weights = INPUT_VARIABLE(1);
|
||||
auto bias = block.width() > 2 ? INPUT_VARIABLE(2) : nullptr;
|
||||
|
||||
auto output = INPUT_VARIABLE(0);
|
||||
|
||||
sd::LongType dD = INT_ARG(9); // dilations depth
|
||||
sd::LongType dH = INT_ARG(10); // dilations height
|
||||
sd::LongType dW = INT_ARG(11); // dilations width
|
||||
int isSameMode = INT_ARG(12); // 0-SAME, 1-VALID
|
||||
|
||||
Requirements req("ONEDNN DECONV3d OP");
|
||||
req.expectTrue(block.isUseONEDNN(), IS_USE_ONEDNN_MSG) &&
|
||||
req.expectLessEq(makeInfoVariable(dD, "Dilation depth"), 1) &&
|
||||
req.expectLessEq(makeInfoVariable(dH, "Dilation height"), 1) &&
|
||||
req.expectLessEq(makeInfoVariable(dW, "Dilation width"), 1) &&
|
||||
req.expectFalse(makeInfoVariable(isSameMode, "isSameMode")) &&
|
||||
req.expectTrue(makeInfoVariable(
|
||||
[input, weights, bias, output] {
|
||||
const DataType xType = input->dataType();
|
||||
const DataType wType = weights->dataType();
|
||||
const DataType zType = output->dataType();
|
||||
const DataType bType = bias != nullptr ? bias->dataType() : zType;
|
||||
return (xType == DataType::FLOAT32 && wType == DataType::FLOAT32 &&
|
||||
bType == DataType::FLOAT32 && zType == DataType::FLOAT32) ||
|
||||
((xType == DataType::UINT8 || xType == DataType::INT8) && wType == DataType::INT8 &&
|
||||
(zType == DataType::UINT8 || zType == DataType::INT8 || zType == DataType::INT32 ||
|
||||
zType == DataType::FLOAT32) &&
|
||||
bType == zType);
|
||||
},
|
||||
TYPECHECK_MSG),
|
||||
NO_MSG);
|
||||
req.logTheSuccess();
|
||||
return req;
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
PLATFORM_IMPL(deconv3d_bp, ENGINE_CPU) {
|
||||
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, oC, iC], [iC, oC, kD, kH, kW], [iC, kD, kH, kW, oC]
|
||||
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), gradI
|
||||
auto gradW = OUTPUT_VARIABLE(1); // [kD, kH, kW, oC, iC], [iC, oC, kD, kH, kW], [iC, kD, kH, kW, oC]
|
||||
auto gradB = block.width() > 3 ? OUTPUT_VARIABLE(2) : nullptr; // [oC]
|
||||
|
||||
REQUIRE_TRUE(input->rankOf() == 5, 0,
|
||||
"CUSTOM DECONV3D_MKLDNN_BP OP: rank of input array must be equal to 5, but got %i instead !",
|
||||
input->rankOf());
|
||||
REQUIRE_TRUE(weights->rankOf() == 5, 0,
|
||||
"CUSTOM DECONV3D_MKLDNN_BP OP: rank of weights array must be equal to 5 , but got %i instead !",
|
||||
weights->rankOf());
|
||||
REQUIRE_TRUE(gradO->rankOf() == 5, 0,
|
||||
"CUSTOM DECONV3D_MKLDNN_BP OP: rank of output gradients (next epsilon) array must be equal to 5, but "
|
||||
"got %i instead !",
|
||||
gradO->rankOf());
|
||||
|
||||
sd::LongType kD = INT_ARG(0) > 0 ? INT_ARG(0) : static_cast<sd::LongType>(weights->sizeAt(0)); // filter(kernel) depth
|
||||
sd::LongType kH = INT_ARG(1) > 0 ? INT_ARG(1) : static_cast<sd::LongType>(weights->sizeAt(1)); // filter(kernel) height
|
||||
sd::LongType kW = INT_ARG(2) > 0 ? INT_ARG(2) : static_cast<sd::LongType>(weights->sizeAt(2)); // filter(kernel) width
|
||||
sd::LongType sD = INT_ARG(3); // strides depth
|
||||
sd::LongType sH = INT_ARG(4); // strides height
|
||||
sd::LongType sW = INT_ARG(5); // strides width
|
||||
sd::LongType pD = INT_ARG(6); // paddings depth
|
||||
sd::LongType pH = INT_ARG(7); // paddings height
|
||||
sd::LongType pW = INT_ARG(8); // paddings width
|
||||
sd::LongType dD = INT_ARG(9); // dilations depth
|
||||
sd::LongType dH = INT_ARG(10); // dilations height
|
||||
sd::LongType dW = INT_ARG(11); // dilations width
|
||||
int isSameMode = 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, oC, iC], 1 - [iC, oC, kD, kH, kW], 2 - [iC, kD, kH, kW, oC]
|
||||
|
||||
sd::LongType bS, iC, iD, iH, iW, oC, oD, oH,
|
||||
oW; // batch size, input channels, input depth/height/width, output channels, output depth/height/width;
|
||||
sd::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, indWoC, indWiC, indWkD);
|
||||
|
||||
sd::LongType trueoD, trueoH, trueoW; // true output height, width
|
||||
ConvolutionUtils::calcOutSizeDeconv3D(trueoD, trueoH, trueoW, kD, kH, kW, sD, sH, sW, pD, pH, pW, dD, dH, dW, iD, iH,
|
||||
iW, isSameMode);
|
||||
|
||||
std::vector<sd::LongType> expectedGradOShape = ShapeUtils::composeShapeUsingDimsAndIdx(
|
||||
{bS, oC, trueoD, trueoH, trueoW, 0, indIOioC, indIOioD, indIOioD + 1, indIOioD + 2});
|
||||
std::vector<sd::LongType> expectedWeightsShape = ConvolutionUtils::expectWeightsShape(wFormat, kD, kH, kW, oC, iC);
|
||||
REQUIRE_TRUE(gradO->isSameShape(expectedGradOShape), 0,
|
||||
"CUSTOM DECONV3D_MKLDNN_BP 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 DECONV3D_MKLDNN_BP 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 DECONV3D_MKLDNN_BP OP: wrong shape of array with biases, expected rank, length: <=2, %i, but "
|
||||
"got %i, %i instead !",
|
||||
oC, bias->rankOf(), bias->lengthOf());
|
||||
|
||||
if (isSameMode) // Note: we're intentionally swapping iH and oH, to calculated the padding for a"normal" conv (not
|
||||
// deconv) forward pass
|
||||
ConvolutionUtils::calcPadding3D(pD, pH, pW, iD, iH, iW, oD, oH, oW, kD, kH, kW, sD, sH, sW, dD, dH, dW);
|
||||
|
||||
deconv3dBackPropMKLDNN(input, weights, gradO, gradI, gradW, gradB, kD, kH, kW, sD, sH, sW, pD, pH, pW, dD, dH, dW,
|
||||
isNCDHW, wFormat);
|
||||
|
||||
return sd::Status::OK;
|
||||
}
|
||||
|
||||
PLATFORM_CHECK(deconv3d_bp, ENGINE_CPU) {
|
||||
auto input = INPUT_VARIABLE(0); // [bS, iD, iH, iW, iC] (NHWC) or [bS, iD, iC, iH, iW] (NCDHW)
|
||||
auto weights = INPUT_VARIABLE(1); // [kD, kH, kW, oC, iC], [iC, oC, kD, kH, kW], [iC, kD, kH, kW, oC]
|
||||
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] (NHWC) or [bS, oC, oD, oH, oW] (NCDHW), epsilon_next
|
||||
|
||||
auto gradI = OUTPUT_VARIABLE(0); // [bS, iD, iH, iW, iC] (NHWC) or [bS, iC, iD, iH, iW] (NCDHW), gradI
|
||||
auto gradW = OUTPUT_VARIABLE(1); // [kD, kH, kW, oC, iC], [iC, oC, kD, kH, kW], [iC, kD, kH, kW, oC]
|
||||
auto gradB = block.width() > 3 ? OUTPUT_VARIABLE(2) : nullptr; // [oC]
|
||||
|
||||
int dD = INT_ARG(9); // dilations depth
|
||||
int dH = INT_ARG(10); // dilations height
|
||||
int dW = INT_ARG(11); // dilations width
|
||||
int isSameMode = INT_ARG(12); // 0-SAME, 1-VALID
|
||||
|
||||
Requirements req("ONEDNN DECONV3d_BP OP");
|
||||
req.expectTrue(block.isUseONEDNN(), IS_USE_ONEDNN_MSG) &&
|
||||
req.expectLessEq(makeInfoVariable(dD, "Dilation depth"), 1) &&
|
||||
req.expectLessEq(makeInfoVariable(dH, "Dilation height"), 1) &&
|
||||
req.expectLessEq(makeInfoVariable(dW, "Dilation width"), 1) &&
|
||||
req.expectFalse(makeInfoVariable(isSameMode, "isSameMode")) &&
|
||||
req.expectTrue(makeInfoVariable(
|
||||
[input, weights, gradO, gradI, gradW, gradB] {
|
||||
const DataType xType = input->dataType();
|
||||
const DataType wType = weights->dataType();
|
||||
const DataType gradOType = gradO->dataType();
|
||||
|
||||
const DataType gradIType = gradI->dataType();
|
||||
const DataType gradWType = gradW->dataType();
|
||||
const DataType gradBType = gradB != nullptr ? gradB->dataType() : DataType::FLOAT32;
|
||||
return ((xType == DataType::FLOAT32 || xType == DataType::BFLOAT16) &&
|
||||
(wType == DataType::FLOAT32 || wType == DataType::BFLOAT16) &&
|
||||
(gradOType == DataType::FLOAT32 || gradOType == DataType::BFLOAT16) &&
|
||||
(gradIType == DataType::FLOAT32 || gradIType == DataType::BFLOAT16) &&
|
||||
(gradWType == DataType::FLOAT32 || gradWType == DataType::BFLOAT16) &&
|
||||
(gradBType == DataType::FLOAT32 || gradBType == DataType::BFLOAT16));
|
||||
},
|
||||
TYPECHECK_MSG),
|
||||
NO_MSG);
|
||||
req.logTheSuccess();
|
||||
return req;
|
||||
}
|
||||
|
||||
} // namespace platforms
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,570 @@
|
||||
/*
|
||||
* ******************************************************************************
|
||||
* *
|
||||
* *
|
||||
* * 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 <helpers/MKLDNNStream.h>
|
||||
#include <ops/declarable/OpRegistrator.h>
|
||||
#include <ops/declarable/PlatformHelper.h>
|
||||
#include <ops/declarable/helpers/convolutions.h>
|
||||
#include <system/platform_boilerplate.h>
|
||||
|
||||
#include "mkldnnUtils.h"
|
||||
|
||||
using namespace dnnl;
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace platforms {
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
static void depthwiseConv2dMKLDNN(NDArray* input, NDArray* weights, NDArray* bias, NDArray* output,
|
||||
const sd::LongType kH, const sd::LongType kW, const sd::LongType sH, const sd::LongType sW, const sd::LongType pH, const sd::LongType pW,
|
||||
const sd::LongType dH, const sd::LongType dW, const int paddingMode, const bool isNCHW,
|
||||
const int wFormat) {
|
||||
// mkl supports only following case: mC = 1, oC = iC
|
||||
|
||||
// input [bS, iC, iH, iW] nchw or [bS, iH, iW, iC] nhwc, since mkl doesn't support nhwc format we'll permute when nhwc
|
||||
// is given weights {iC, mC, 1, kH, kW} bias [oC], may be nullptr output [bS, oC, oH, oW] nchw or [bS, oH, oW, oC]
|
||||
// nhwc oC = iC*mC
|
||||
|
||||
sd::LongType bS, iC, iH, iW, mC, oC, oH,
|
||||
oW; // batch size, input channels, input height/width, output channels, output height/width;
|
||||
sd::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
|
||||
|
||||
const int pWSame = (paddingMode == 2 && dW > 1) ? ((oW - 1) * sW + (kW - 1) * dW + 1 - iW) / 2
|
||||
: pW; // dH == 1 for causal mode in conv1d
|
||||
|
||||
dnnl::memory::dims strides = {sH, sW};
|
||||
dnnl::memory::dims padding = {pH, pW};
|
||||
dnnl::memory::dims padding_r = {(oH - 1) * sH - iH + kH - pH, (oW - 1) * sW - iW + kW - pWSame};
|
||||
dnnl::memory::dims dilation = {dH - 1, dW - 1};
|
||||
|
||||
sd::LongType i0, i1, i2, i3;
|
||||
if (0 == wFormat) {
|
||||
i0 = 2;
|
||||
i1 = 3;
|
||||
i2 = 0;
|
||||
i3 = 1; // [kH, kW, iC, mC] -> [iC, mC, 1, kH, kW]
|
||||
} else if (1 == wFormat) {
|
||||
i0 = 1;
|
||||
i1 = 0;
|
||||
i2 = 2;
|
||||
i3 = 3; // [mC, iC, kH, kW] -> [iC, mC, 1, kH, kW]
|
||||
} else {
|
||||
i0 = 3;
|
||||
i1 = 0;
|
||||
i2 = 1;
|
||||
i3 = 2; // [mC, kH, kW, iC] -> [iC, mC, 1, kH, kW]
|
||||
}
|
||||
|
||||
// input type
|
||||
dnnl::memory::data_type xType;
|
||||
if (input->dataType() == DataType::FLOAT32)
|
||||
xType = dnnl::memory::data_type::f32;
|
||||
else if (input->dataType() == DataType::HALF)
|
||||
xType = dnnl::memory::data_type::f16;
|
||||
else if (input->dataType() == DataType::UINT8)
|
||||
xType = dnnl::memory::data_type::u8;
|
||||
else
|
||||
xType = dnnl::memory::data_type::s8;
|
||||
|
||||
// weights type
|
||||
dnnl::memory::data_type wType = xType;
|
||||
if (xType == dnnl::memory::data_type::u8) wType = dnnl::memory::data_type::s8;
|
||||
|
||||
// output and bias type (have the same types)
|
||||
dnnl::memory::data_type zType;
|
||||
if (output->dataType() == DataType::FLOAT32)
|
||||
zType = dnnl::memory::data_type::f32;
|
||||
else if (output->dataType() == DataType::HALF)
|
||||
zType = dnnl::memory::data_type::f16;
|
||||
else if (output->dataType() == DataType::UINT8)
|
||||
zType = dnnl::memory::data_type::u8;
|
||||
else if (output->dataType() == DataType::INT8)
|
||||
zType = dnnl::memory::data_type::s8;
|
||||
else
|
||||
zType = dnnl::memory::data_type::s32;
|
||||
|
||||
dnnl::memory::format_tag xzFormatMkl = isNCHW ? dnnl::memory::format_tag::nchw : dnnl::memory::format_tag::nhwc;
|
||||
dnnl::memory::format_tag wFormatMkl = dnnl::memory::format_tag::goihw;
|
||||
|
||||
dnnl::memory::dims xDims = {bS, iC, iH, iW};
|
||||
dnnl::memory::dims wDims = {iC, mC, 1, kH, kW};
|
||||
dnnl::memory::dims zDims = {bS, oC, oH, oW};
|
||||
|
||||
// memory descriptors for arrays
|
||||
|
||||
// input
|
||||
dnnl::memory::desc x_mkl_md = dnnl::memory::desc(xDims, xType, dnnl::memory::format_tag::any);
|
||||
dnnl::memory::desc x_user_md = dnnl::memory::desc(xDims, xType, xzFormatMkl);
|
||||
onednnUtils::setBlockStrides(*input, x_user_md);
|
||||
|
||||
// weights
|
||||
dnnl::memory::desc w_mkl_md = dnnl::memory::desc(wDims, wType, dnnl::memory::format_tag::any);
|
||||
dnnl::memory::desc w_user_md = dnnl::memory::desc(wDims, wType, wFormatMkl);
|
||||
w_user_md.data.format_kind = dnnl_blocked; // overrides format
|
||||
w_user_md.data.format_desc.blocking.strides[0] = weights->strideAt(i0); // permute
|
||||
w_user_md.data.format_desc.blocking.strides[1] = weights->strideAt(i1);
|
||||
w_user_md.data.format_desc.blocking.strides[2] = 0;
|
||||
w_user_md.data.format_desc.blocking.strides[3] = weights->strideAt(i2);
|
||||
w_user_md.data.format_desc.blocking.strides[4] = weights->strideAt(i3);
|
||||
|
||||
// bias
|
||||
dnnl::memory::desc b_mkl_md;
|
||||
if (bias != nullptr) b_mkl_md = dnnl::memory::desc({oC}, zType, dnnl::memory::format_tag::x);
|
||||
|
||||
// output
|
||||
dnnl::memory::desc z_mkl_md = dnnl::memory::desc(zDims, zType, dnnl::memory::format_tag::any);
|
||||
dnnl::memory::desc z_user_md = dnnl::memory::desc(zDims, zType, xzFormatMkl);
|
||||
onednnUtils::setBlockStrides(*output, z_user_md);
|
||||
|
||||
auto engine = onednnUtils::getEngine(LaunchContext::defaultContext()->engine());
|
||||
|
||||
// operation primitive description
|
||||
dnnl::convolution_forward::desc op_desc(dnnl::prop_kind::forward_inference, dnnl::algorithm::convolution_auto,
|
||||
x_mkl_md, w_mkl_md, b_mkl_md, z_mkl_md, strides, dilation, padding,
|
||||
padding_r);
|
||||
dnnl::convolution_forward::primitive_desc op_prim_desc(op_desc, engine);
|
||||
|
||||
// arguments (memory buffers) necessary for calculations
|
||||
std::unordered_map<int, dnnl::memory> args;
|
||||
|
||||
dnnl::stream stream(engine);
|
||||
|
||||
// provide memory buffers and check whether reorder is required
|
||||
|
||||
// input
|
||||
onednnUtils::loadDataToMklStream(*input, engine, stream, x_user_md, op_prim_desc.src_desc(), args[DNNL_ARG_SRC]);
|
||||
|
||||
// weights
|
||||
onednnUtils::loadDataToMklStream(*weights, engine, stream, w_user_md, op_prim_desc.weights_desc(),
|
||||
args[DNNL_ARG_WEIGHTS]);
|
||||
|
||||
// bias
|
||||
if (bias != nullptr) {
|
||||
auto b_mkl_mem = dnnl::memory(b_mkl_md, engine, const_cast<void*>(bias->buffer()));
|
||||
args[DNNL_ARG_BIAS] = b_mkl_mem;
|
||||
}
|
||||
|
||||
// output
|
||||
auto z_user_mem =
|
||||
onednnUtils::loadDataToMklStream(*output, engine, stream, z_user_md, op_prim_desc.dst_desc(), args[DNNL_ARG_DST]);
|
||||
|
||||
// run calculations
|
||||
dnnl::convolution_forward(op_prim_desc).execute(stream, args);
|
||||
|
||||
// reorder outputs if necessary
|
||||
if (op_prim_desc.dst_desc() != z_user_mem.get_desc())
|
||||
dnnl::reorder(args[DNNL_ARG_DST], z_user_mem).execute(stream, args[DNNL_ARG_DST], z_user_mem);
|
||||
|
||||
stream.wait();
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
static void depthwiseConv2dBpMKLDNN(NDArray* input, NDArray* weights, NDArray* gradO, NDArray* gradI,
|
||||
NDArray* gradW, NDArray* gradB, const sd::LongType kH, const sd::LongType kW, const sd::LongType sH,
|
||||
const sd::LongType sW, const sd::LongType pH, const sd::LongType pW, const sd::LongType dH, const sd::LongType dW,
|
||||
const int paddingMode, const bool isNCHW, const int wFormat) {
|
||||
// mkl supports only following case: mC = 1, oC = iC
|
||||
|
||||
// input, gradI [bS, iC, iH, iW] nchw or [bS, iH, iW, iC] nhwc, since mkl doesn't support nhwc format we'll permute
|
||||
// when nhwc is given weights/gradW {iC, mC, 1, kH, kW} gradB [oC], may be nullptr gradO [bS, oC, oH, oW] nchw or [bS,
|
||||
// oH, oW, oC] nhwc oC = iC*mC
|
||||
|
||||
sd::LongType bS, iC, iH, iW, mC, oC, oH,
|
||||
oW; // batch size, input channels, input height/width, output channels, output height/width;
|
||||
sd::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);
|
||||
|
||||
const int pWSame = (paddingMode == 2 && dW > 1) ? ((oW - 1) * sW + (kW - 1) * dW + 1 - iW) / 2
|
||||
: pW; // dH == 1 for causal mode in conv1d
|
||||
|
||||
dnnl::memory::dims strides = {sH, sW};
|
||||
dnnl::memory::dims padding = {pH, pW};
|
||||
dnnl::memory::dims padding_r = {(oH - 1) * sH - iH + kH - pH, (oW - 1) * sW - iW + kW - pWSame};
|
||||
dnnl::memory::dims dilation = {dH - 1, dW - 1};
|
||||
|
||||
sd::LongType i0, i1, i2, i3;
|
||||
if (0 == wFormat) {
|
||||
i0 = 2;
|
||||
i1 = 3;
|
||||
i2 = 0;
|
||||
i3 = 1; // [kH, kW, iC, mC] -> [iC, mC, 1, kH, kW]
|
||||
} else if (1 == wFormat) {
|
||||
i0 = 1;
|
||||
i1 = 0;
|
||||
i2 = 2;
|
||||
i3 = 3; // [mC, iC, kH, kW] -> [iC, mC, 1, kH, kW]
|
||||
} else {
|
||||
i0 = 3;
|
||||
i1 = 0;
|
||||
i2 = 1;
|
||||
i3 = 2; // [mC, kH, kW, iC] -> [iC, mC, 1, kH, kW]
|
||||
}
|
||||
|
||||
// input type
|
||||
dnnl::memory::data_type xType =
|
||||
input->dataType() == DataType::FLOAT32 ? dnnl::memory::data_type::f32 : dnnl::memory::data_type::bf16;
|
||||
// weights type
|
||||
dnnl::memory::data_type wType =
|
||||
weights->dataType() == DataType::FLOAT32 ? dnnl::memory::data_type::f32 : dnnl::memory::data_type::bf16;
|
||||
// gradO type
|
||||
dnnl::memory::data_type gradOType =
|
||||
gradO->dataType() == DataType::FLOAT32 ? dnnl::memory::data_type::f32 : dnnl::memory::data_type::bf16;
|
||||
// gradI type
|
||||
dnnl::memory::data_type gradIType =
|
||||
gradI->dataType() == DataType::FLOAT32 ? dnnl::memory::data_type::f32 : dnnl::memory::data_type::bf16;
|
||||
// gradW type
|
||||
dnnl::memory::data_type gradWType =
|
||||
gradW->dataType() == DataType::FLOAT32 ? dnnl::memory::data_type::f32 : dnnl::memory::data_type::bf16;
|
||||
// gradB type
|
||||
dnnl::memory::data_type gradBType =
|
||||
gradB != nullptr
|
||||
? (gradB->dataType() == DataType::FLOAT32 ? dnnl::memory::data_type::f32 : dnnl::memory::data_type::bf16)
|
||||
: dnnl::memory::data_type::f32;
|
||||
|
||||
dnnl::memory::format_tag xzFormatMkl = isNCHW ? dnnl::memory::format_tag::nchw : dnnl::memory::format_tag::nhwc;
|
||||
dnnl::memory::format_tag wFormatMkl = dnnl::memory::format_tag::goihw;
|
||||
|
||||
dnnl::memory::dims xDims = {bS, iC, iH, iW};
|
||||
dnnl::memory::dims wDims = {iC, mC, 1, kH, kW};
|
||||
dnnl::memory::dims zDims = {bS, oC, oH, oW};
|
||||
|
||||
// memory descriptors for arrays
|
||||
|
||||
// input
|
||||
dnnl::memory::desc x_mkl_md = dnnl::memory::desc(xDims, xType, dnnl::memory::format_tag::any);
|
||||
dnnl::memory::desc x_user_md = dnnl::memory::desc(xDims, xType, xzFormatMkl);
|
||||
onednnUtils::setBlockStrides(*input, x_user_md);
|
||||
|
||||
// weights
|
||||
dnnl::memory::desc w_mkl_md = dnnl::memory::desc(wDims, wType, dnnl::memory::format_tag::any);
|
||||
dnnl::memory::desc w_user_md = dnnl::memory::desc(wDims, wType, wFormatMkl);
|
||||
w_user_md.data.format_kind = dnnl_blocked; // overrides format
|
||||
w_user_md.data.format_desc.blocking.strides[0] = weights->strideAt(i0); // permute
|
||||
w_user_md.data.format_desc.blocking.strides[1] = weights->strideAt(i1);
|
||||
w_user_md.data.format_desc.blocking.strides[2] = 0;
|
||||
w_user_md.data.format_desc.blocking.strides[3] = weights->strideAt(i2);
|
||||
w_user_md.data.format_desc.blocking.strides[4] = weights->strideAt(i3);
|
||||
|
||||
// gradO
|
||||
dnnl::memory::desc gradO_mkl_md = dnnl::memory::desc(zDims, gradOType, dnnl::memory::format_tag::any);
|
||||
dnnl::memory::desc gradO_user_md = dnnl::memory::desc(zDims, gradOType, xzFormatMkl);
|
||||
onednnUtils::setBlockStrides(*gradO, gradO_user_md);
|
||||
|
||||
// gradI
|
||||
dnnl::memory::desc gradI_mkl_md = dnnl::memory::desc(xDims, gradIType, dnnl::memory::format_tag::any);
|
||||
dnnl::memory::desc gradI_user_md = dnnl::memory::desc(xDims, gradIType, xzFormatMkl);
|
||||
onednnUtils::setBlockStrides(*gradI, gradI_user_md);
|
||||
|
||||
// gradW
|
||||
dnnl::memory::desc gradW_mkl_md = dnnl::memory::desc(wDims, gradWType, dnnl::memory::format_tag::any);
|
||||
dnnl::memory::desc gradW_user_md = dnnl::memory::desc(wDims, gradWType, wFormatMkl);
|
||||
gradW_user_md.data.format_kind = dnnl_blocked; // overrides format
|
||||
gradW_user_md.data.format_desc.blocking.strides[0] = gradW->strideAt(i0); // permute
|
||||
gradW_user_md.data.format_desc.blocking.strides[1] = gradW->strideAt(i1);
|
||||
gradW_user_md.data.format_desc.blocking.strides[2] = 0;
|
||||
gradW_user_md.data.format_desc.blocking.strides[3] = gradW->strideAt(i2);
|
||||
gradW_user_md.data.format_desc.blocking.strides[4] = gradW->strideAt(i3);
|
||||
|
||||
// gradB
|
||||
dnnl::memory::desc gradB_mkl_md;
|
||||
if (gradB != nullptr) gradB_mkl_md = dnnl::memory::desc({oC}, gradBType, dnnl::memory::format_tag::x);
|
||||
|
||||
auto engine = onednnUtils::getEngine(LaunchContext::defaultContext()->engine());
|
||||
|
||||
// forward primitive description
|
||||
dnnl::convolution_forward::desc op_ff_desc(dnnl::prop_kind::forward_inference, dnnl::algorithm::convolution_auto,
|
||||
x_mkl_md, w_mkl_md, gradB_mkl_md, gradO_mkl_md, strides, dilation, padding,
|
||||
padding_r);
|
||||
dnnl::convolution_forward::primitive_desc op_ff_prim_desc(op_ff_desc, engine);
|
||||
|
||||
// backward data primitive description
|
||||
dnnl::convolution_backward_data::desc op_data_bp_desc(dnnl::algorithm::convolution_auto, gradI_mkl_md, w_mkl_md,
|
||||
gradO_mkl_md, strides, dilation, padding, padding_r);
|
||||
dnnl::convolution_backward_data::primitive_desc op_data_bp_prim_desc(op_data_bp_desc, engine, op_ff_prim_desc);
|
||||
|
||||
// backward weights primitive description
|
||||
dnnl::convolution_backward_weights::desc op_weights_bp_desc(dnnl::algorithm::convolution_auto, x_mkl_md, gradW_mkl_md,
|
||||
gradB_mkl_md, gradO_mkl_md, strides, dilation, padding,
|
||||
padding_r);
|
||||
dnnl::convolution_backward_weights::primitive_desc op_weights_bp_prim_desc(op_weights_bp_desc, engine,
|
||||
op_ff_prim_desc);
|
||||
|
||||
// arguments (memory buffers) necessary for calculations
|
||||
std::unordered_map<int, dnnl::memory> args;
|
||||
|
||||
dnnl::stream stream(engine);
|
||||
|
||||
// provide memory buffers and check whether reorder is required
|
||||
|
||||
// input
|
||||
onednnUtils::loadDataToMklStream(*input, engine, stream, x_user_md, op_weights_bp_prim_desc.src_desc(),
|
||||
args[DNNL_ARG_SRC]);
|
||||
|
||||
// weights
|
||||
onednnUtils::loadDataToMklStream(*weights, engine, stream, w_user_md, op_data_bp_prim_desc.weights_desc(),
|
||||
args[DNNL_ARG_WEIGHTS]);
|
||||
|
||||
// gradO
|
||||
auto gradO_user_mem = dnnl::memory(gradO_user_md, engine, const_cast<void*>(gradO->buffer()));
|
||||
const bool gradOReorderW = op_weights_bp_prim_desc.diff_dst_desc() != gradO_user_mem.get_desc();
|
||||
const bool gradOReorderD = op_data_bp_prim_desc.diff_dst_desc() != gradO_user_mem.get_desc();
|
||||
auto gradO_mkl_memW = gradOReorderW ? dnnl::memory(op_weights_bp_prim_desc.diff_dst_desc(), engine) : gradO_user_mem;
|
||||
auto gradO_mkl_memD = gradOReorderD ? dnnl::memory(op_data_bp_prim_desc.diff_dst_desc(), engine) : gradO_user_mem;
|
||||
if (gradOReorderW) dnnl::reorder(gradO_user_mem, gradO_mkl_memW).execute(stream, gradO_user_mem, gradO_mkl_memW);
|
||||
if (gradOReorderD) dnnl::reorder(gradO_user_mem, gradO_mkl_memD).execute(stream, gradO_user_mem, gradO_mkl_memD);
|
||||
args[DNNL_ARG_DIFF_DST] = gradO_mkl_memD;
|
||||
|
||||
// gradI
|
||||
auto gradI_user_mem = onednnUtils::loadDataToMklStream(*gradI, engine, stream, gradI_user_md,
|
||||
op_data_bp_prim_desc.diff_src_desc(), args[DNNL_ARG_DIFF_SRC]);
|
||||
|
||||
// gradW
|
||||
auto gradW_user_mem = onednnUtils::loadDataToMklStream(
|
||||
*gradW, engine, stream, gradW_user_md, op_weights_bp_prim_desc.diff_weights_desc(), args[DNNL_ARG_DIFF_WEIGHTS]);
|
||||
|
||||
// gradB
|
||||
if (gradB != nullptr) {
|
||||
auto gradB_mkl_mem = dnnl::memory(gradB_mkl_md, engine, gradB->buffer());
|
||||
args[DNNL_ARG_DIFF_BIAS] = gradB_mkl_mem;
|
||||
}
|
||||
|
||||
// run backward data calculations
|
||||
dnnl::convolution_backward_data(op_data_bp_prim_desc).execute(stream, args);
|
||||
|
||||
if (gradOReorderW || gradOReorderD) args[DNNL_ARG_DIFF_DST] = gradO_mkl_memW;
|
||||
|
||||
// run backward weights calculations
|
||||
dnnl::convolution_backward_weights(op_weights_bp_prim_desc).execute(stream, args);
|
||||
|
||||
// reorder gradI if necessary
|
||||
if (op_data_bp_prim_desc.diff_src_desc() != gradI_user_mem.get_desc())
|
||||
dnnl::reorder(args[DNNL_ARG_DIFF_SRC], gradI_user_mem).execute(stream, args[DNNL_ARG_DIFF_SRC], gradI_user_mem);
|
||||
if (op_weights_bp_prim_desc.diff_weights_desc() != gradW_user_mem.get_desc())
|
||||
dnnl::reorder(args[DNNL_ARG_DIFF_WEIGHTS], gradW_user_mem)
|
||||
.execute(stream, args[DNNL_ARG_DIFF_WEIGHTS], gradW_user_mem);
|
||||
|
||||
stream.wait();
|
||||
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////
|
||||
PLATFORM_IMPL(depthwise_conv2d, ENGINE_CPU) {
|
||||
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)
|
||||
|
||||
sd::LongType kH = INT_ARG(0) > 0 ? INT_ARG(0) : static_cast<sd::LongType>(weights->sizeAt(0)); // filter(kernel) height
|
||||
sd::LongType kW = INT_ARG(1) > 0 ? INT_ARG(1) : static_cast<sd::LongType>(weights->sizeAt(1)); // filter(kernel) width
|
||||
sd::LongType sH = INT_ARG(2); // strides height
|
||||
sd::LongType sW = INT_ARG(3); // strides width
|
||||
sd::LongType pH = INT_ARG(4); // paddings height
|
||||
sd::LongType pW = INT_ARG(5); // paddings width
|
||||
sd::LongType dH = INT_ARG(6); // dilations height
|
||||
sd::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]
|
||||
|
||||
sd::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
|
||||
sd::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<sd::LongType> expectedWeightsShape = ConvolutionUtils::expectWeightsShape(wFormat, kH, kW, iC, mC);
|
||||
REQUIRE_TRUE(weights->isSameShape(expectedWeightsShape), 0,
|
||||
"CUSTOM DEPTHWISECONV2D MKL 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,
|
||||
"CUSTOM DEPTHWISECONV2D MKL 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,
|
||||
"CUSTOM DEPTHWISECONV2D MKL OP: wrong shape of array with biases, expected rank, length: <=2, %i, but "
|
||||
"got %i, %i instead !",
|
||||
oC, bias->rankOf(), bias->lengthOf());
|
||||
|
||||
depthwiseConv2dMKLDNN(input, weights, bias, output, kH, kW, sH, sW, pH, pW, dH, dW, paddingMode, isNCHW, wFormat);
|
||||
|
||||
return sd::Status::OK;
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////
|
||||
PLATFORM_CHECK(depthwise_conv2d, ENGINE_CPU) {
|
||||
auto input = INPUT_VARIABLE(0);
|
||||
auto weights = INPUT_VARIABLE(1);
|
||||
auto bias = block.width() > 2 ? INPUT_VARIABLE(2) : nullptr;
|
||||
|
||||
auto output = INPUT_VARIABLE(0);
|
||||
|
||||
Requirements req("ONEDNN DEPTHWISE_CONV2d OP");
|
||||
req.expectTrue(block.isUseONEDNN(), IS_USE_ONEDNN_MSG) &&
|
||||
req.expectEq(makeInfoVariable(weights->sizeAt(3), "weight NdArray size#3"), 1) &&
|
||||
req.expectTrue(makeInfoVariable(
|
||||
[input, weights, bias, output] {
|
||||
const DataType xType = input->dataType();
|
||||
const DataType wType = weights->dataType();
|
||||
const DataType zType = output->dataType();
|
||||
const DataType bType = bias != nullptr ? bias->dataType() : zType;
|
||||
return ((xType == DataType::FLOAT32 && wType == DataType::FLOAT32 &&
|
||||
bType == DataType::FLOAT32 && zType == DataType::FLOAT32) ||
|
||||
(xType == DataType::BFLOAT16 && wType == DataType::BFLOAT16 &&
|
||||
bType == DataType::BFLOAT16 && zType == DataType::BFLOAT16) ||
|
||||
((xType == DataType::UINT8 || xType == DataType::INT8) && wType == DataType::INT8 &&
|
||||
(zType == DataType::UINT8 || zType == DataType::INT8 || zType == DataType::INT32 ||
|
||||
zType == DataType::FLOAT32) &&
|
||||
bType == zType));
|
||||
},
|
||||
TYPECHECK_MSG),
|
||||
NO_MSG);
|
||||
req.logTheSuccess();
|
||||
return req;
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
PLATFORM_IMPL(depthwise_conv2d_bp, ENGINE_CPU) {
|
||||
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_NULLIFIED(0); // [bS, iH, iW, iC] (NDHWC) or [bS, iC, iH, iW] (NCDHW), epsilon
|
||||
auto gradW = OUTPUT_NULLIFIED(1); // [kH, kW, iC, mC], [mC, iC, kH, kW], [mC, kH, kW, iC]
|
||||
auto gradB = block.width() > 3 ? OUTPUT_NULLIFIED(2) : nullptr; // [oC]
|
||||
|
||||
REQUIRE_TRUE(input->rankOf() == 4, 0,
|
||||
"CUSTOM DEPTHWISECONV2D_BP MKL OP: rank of input array must be equal to 4, but got %i instead !",
|
||||
input->rankOf());
|
||||
REQUIRE_TRUE(weights->rankOf() == 4, 0,
|
||||
"CUSTOM DEPTHWISECONV2D_BP MKL OP: rank of weights array must be equal to 4, but got %i instead !",
|
||||
weights->rankOf());
|
||||
REQUIRE_TRUE(gradO->rankOf() == 4, 0,
|
||||
"CUSTOM DEPTHWISECONV2D_BP MKL OP: rank of output gradients (next epsilon) array must be equal to 4, "
|
||||
"but got %i instead !",
|
||||
gradO->rankOf());
|
||||
|
||||
sd::LongType kH = INT_ARG(0) > 0 ? INT_ARG(0) : static_cast<sd::LongType>(weights->sizeAt(0)); // filter(kernel) height
|
||||
sd::LongType kW = INT_ARG(1) > 0 ? INT_ARG(1) : static_cast<sd::LongType>(weights->sizeAt(1)); // filter(kernel) width
|
||||
sd::LongType sH = INT_ARG(2); // strides height
|
||||
sd::LongType sW = INT_ARG(3); // strides width
|
||||
sd::LongType pH = INT_ARG(4); // paddings height
|
||||
sd::LongType pW = INT_ARG(5); // paddings width
|
||||
sd::LongType dH = INT_ARG(6); // dilations height
|
||||
sd::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]
|
||||
|
||||
sd::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
|
||||
sd::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
|
||||
|
||||
sd::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<sd::LongType> expectedGradOShape =
|
||||
ShapeUtils::composeShapeUsingDimsAndIdx({bS, oC, trueoH, trueoW, 0, indIOioC, indOoH, indOoH + 1});
|
||||
std::vector<sd::LongType> expectedWeightsShape = ConvolutionUtils::expectWeightsShape(wFormat, kH, kW, iC, mC);
|
||||
REQUIRE_TRUE(gradO->isSameShape(expectedGradOShape), 0,
|
||||
"CUSTOM DEPTHWISECONV2D_BP MKL 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 DEPTHWISECONV2D_BP MKL 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 DEPTHWISECONV2D_BP MKL OP: wrong shape of array with biases, expected rank, length: <=2, %i, "
|
||||
"but got %i, %i instead !",
|
||||
oC, bias->rankOf(), bias->lengthOf());
|
||||
|
||||
depthwiseConv2dBpMKLDNN(input, weights, gradO, gradI, gradW, gradB, kH, kW, sH, sW, pH, pW, dH, dW, paddingMode,
|
||||
isNCHW, wFormat);
|
||||
|
||||
return sd::Status::OK;
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////
|
||||
PLATFORM_CHECK(depthwise_conv2d_bp, ENGINE_CPU) {
|
||||
auto input = INPUT_VARIABLE(0); // [bS, iH, iW, iC] (NDHWC) or [bS, iC, iH, iW] (NCDHW)
|
||||
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] = [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, oC], [oC, iC, kH, kW], [oC, kH, kW, iC]
|
||||
auto gradB = block.width() > 3 ? OUTPUT_VARIABLE(2) : nullptr; // [oC]
|
||||
|
||||
Requirements req("ONEDNN DEPTHWISE_CONV2d_BP OP");
|
||||
req.expectTrue(block.isUseONEDNN(), IS_USE_ONEDNN_MSG) &&
|
||||
req.expectEq(makeInfoVariable(weights->sizeAt(3), "weight NdArray size#3"), 1) &&
|
||||
req.expectTrue(makeInfoVariable(
|
||||
[input, weights, gradI, gradW, gradB, gradO] {
|
||||
const DataType xType = input->dataType();
|
||||
const DataType wType = weights->dataType();
|
||||
const DataType gradOType = gradO->dataType();
|
||||
|
||||
const DataType gradIType = gradI->dataType();
|
||||
const DataType gradWType = gradW->dataType();
|
||||
const DataType gradBType = gradB != nullptr ? gradB->dataType() : DataType::FLOAT32;
|
||||
return ((xType == DataType::FLOAT32 || xType == DataType::BFLOAT16) &&
|
||||
(wType == DataType::FLOAT32 || wType == DataType::BFLOAT16) &&
|
||||
(gradOType == DataType::FLOAT32 || gradOType == DataType::BFLOAT16) &&
|
||||
(gradIType == DataType::FLOAT32 || gradIType == DataType::BFLOAT16) &&
|
||||
(gradWType == DataType::FLOAT32 || gradWType == DataType::BFLOAT16) &&
|
||||
(gradBType == DataType::FLOAT32 || gradBType == DataType::BFLOAT16));
|
||||
},
|
||||
TYPECHECK_MSG),
|
||||
NO_MSG);
|
||||
req.logTheSuccess();
|
||||
return req;
|
||||
}
|
||||
|
||||
} // namespace platforms
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,96 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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 saudet
|
||||
// @author raver119@gmail.com
|
||||
//
|
||||
#include <helpers/MKLDNNStream.h>
|
||||
#include <ops/declarable/OpRegistrator.h>
|
||||
#include <ops/declarable/PlatformHelper.h>
|
||||
#include <ops/declarable/helpers/convolutions.h>
|
||||
#include <system/platform_boilerplate.h>
|
||||
|
||||
#include "mkldnnUtils.h"
|
||||
|
||||
using namespace dnnl;
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace platforms {
|
||||
PLATFORM_IMPL(lrn, ENGINE_CPU) {
|
||||
auto input = INPUT_VARIABLE(0);
|
||||
auto output = OUTPUT_VARIABLE(0);
|
||||
|
||||
REQUIRE_TRUE(input->rankOf() == 4, 0, "lrn: Input rank of 4 expected, but got %i instead", input->rankOf());
|
||||
|
||||
double alpha = T_ARG(1);
|
||||
double beta = T_ARG(2);
|
||||
double bias = T_ARG(0);
|
||||
int depth = INT_ARG(0);
|
||||
|
||||
dnnl_memory_desc_t empty;
|
||||
dnnl::memory::desc lrn_src_md(empty), lrn_dst_md(empty), user_src_md(empty), user_dst_md(empty);
|
||||
|
||||
onednnUtils::getONEDNNMemoryDescLrn(input, nullptr, output, &lrn_src_md, nullptr, &lrn_dst_md, &user_src_md, nullptr,
|
||||
&user_dst_md, input->rankOf() - 1);
|
||||
|
||||
auto lrn_desc = lrn_forward::desc(prop_kind::forward_inference, algorithm::lrn_across_channels, lrn_src_md,
|
||||
(2 * depth + 1), alpha * (2 * depth + 1), beta, bias);
|
||||
|
||||
auto engine = onednnUtils::getEngine(LaunchContext::defaultContext()->engine());
|
||||
dnnl::stream stream(engine);
|
||||
auto lrn_prim_desc = lrn_forward::primitive_desc(lrn_desc, engine);
|
||||
auto user_src_memory = dnnl::memory(user_src_md, engine, input->buffer());
|
||||
auto user_dst_memory = dnnl::memory(user_dst_md, engine, output->buffer());
|
||||
|
||||
auto lrn_src_memory = user_src_memory;
|
||||
if (lrn_prim_desc.src_desc() != user_src_memory.get_desc()) {
|
||||
lrn_src_memory = dnnl::memory(lrn_prim_desc.src_desc(), engine);
|
||||
reorder(user_src_memory, lrn_src_memory).execute(stream, user_src_memory, lrn_src_memory);
|
||||
}
|
||||
|
||||
auto lrn_dst_memory = user_dst_memory;
|
||||
if (lrn_prim_desc.dst_desc() != user_dst_memory.get_desc()) {
|
||||
lrn_dst_memory = dnnl::memory(lrn_prim_desc.dst_desc(), engine);
|
||||
}
|
||||
|
||||
lrn_forward(lrn_prim_desc).execute(stream, {{DNNL_ARG_SRC, lrn_src_memory}, {DNNL_ARG_DST, lrn_dst_memory}});
|
||||
|
||||
if (lrn_prim_desc.dst_desc() != user_dst_memory.get_desc()) {
|
||||
reorder(lrn_dst_memory, user_dst_memory).execute(stream, lrn_dst_memory, user_dst_memory);
|
||||
}
|
||||
|
||||
stream.wait();
|
||||
|
||||
return sd::Status::OK;
|
||||
};
|
||||
|
||||
PLATFORM_CHECK(lrn, ENGINE_CPU) {
|
||||
auto input = INPUT_VARIABLE(0);
|
||||
auto output = OUTPUT_VARIABLE(0);
|
||||
|
||||
Requirements req("ONEDNN LRN OP");
|
||||
req.expectTrue(block.isUseONEDNN(), IS_USE_ONEDNN_MSG) &&
|
||||
req.expectTrue(sd::ONEDNNStream::isSupported({input, output}), ONEDNN_STREAM_NOT_SUPPORTED);
|
||||
req.logTheSuccess();
|
||||
return req;
|
||||
}
|
||||
} // namespace platforms
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,547 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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/OpRegistrator.h>
|
||||
|
||||
#include "mkldnnUtils.h"
|
||||
|
||||
using namespace dnnl;
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace platforms {
|
||||
|
||||
static void lstmLayerMKLDNN(NDArray* x, NDArray* Wx, NDArray* Wr, NDArray* b, NDArray* hI,
|
||||
NDArray* cI, const std::vector<float>& params, NDArray* h, NDArray* hL, NDArray* cL) {
|
||||
// equations (no peephole connections)
|
||||
// it = σ(Wxi * xt + Wri * ht-1 + bi)
|
||||
// ft = σ(Wxf * xt + Wrf * ht-1 + bf)
|
||||
// c't = tanh(Wxc * xt + Wrc * ht-1 + bc)
|
||||
// ct = ft ◦ ct-1 + it ◦ c't
|
||||
// ot = σ(Wxo * xt + Wro * ht-1 + bo)
|
||||
// ht = ot ◦ tanh(ct)
|
||||
|
||||
// notations:
|
||||
// bS - batch size
|
||||
// sL - sequence length, number of time steps
|
||||
// nIn - input size
|
||||
// nOut - output size (hidden size)
|
||||
|
||||
// INPUTS:
|
||||
|
||||
// *******
|
||||
// input x:
|
||||
// 1) [sL, bS, nIn] when dataFormat == 0
|
||||
|
||||
// *******
|
||||
// input weights Wx:
|
||||
// 1) [1, 1, nIn, 4*nOut] when directionMode < 2
|
||||
// 2) [1, 2, nIn, 4*nOut] when directionMode >= 2
|
||||
|
||||
// *******
|
||||
// recurrent weights Wr:
|
||||
// 1) [1, 1, nOut, 4*nOut] when directionMode < 2
|
||||
// 2) [1, 2, nOut, 4*nOut] when directionMode >= 2
|
||||
|
||||
// *******
|
||||
// biases b:
|
||||
// 1) [1, 1, 4*nOut] when directionMode < 2
|
||||
// 2) [1, 2, 4*nOut] when directionMode >= 2
|
||||
|
||||
// *******
|
||||
// initial output hI:
|
||||
// 1) [1, 1, bS, nOut] when directionMode < 2
|
||||
// 2) [1, 2, bS, nOut] when directionMode >= 2
|
||||
|
||||
// *******
|
||||
// initial cell state cI (same shape as in hI):
|
||||
// 1) [1, 1, bS, nOut] when directionMode < 2
|
||||
// 2) [1, 2, bS, nOut] when directionMode >= 2
|
||||
|
||||
// OUTPUTS:
|
||||
|
||||
// *******
|
||||
// output h:
|
||||
// 1) [sL, bS, nOut] when directionMode <= 2 && dataFormat == 0
|
||||
// 2) [sL, bS, 2*nOut] when directionMode == 3 && dataFormat == 0
|
||||
|
||||
// *******
|
||||
// output at last step hL:
|
||||
// 1) [1, 1, bS, nOut] when directionMode < 2
|
||||
// 2) [1, 2, bS, nOut] when directionMode >= 2
|
||||
|
||||
// *******
|
||||
// cell state at last step cL (same shape as in hL):
|
||||
// 1) [1, 1, bS, nOut] when directionMode < 2
|
||||
// 2) [1, 2, bS, nOut] when directionMode >= 2
|
||||
|
||||
// !!! dimension 4*nOut implies order it, ft, c't, ot
|
||||
// !!! dimension 3*nOut implies order it, ft, ot
|
||||
|
||||
// params = {dataFormat, directionMode, cellClip, gateAct, gateAlpha, gateBeta, cellAct, cellAlpha, cellBeta, outAct,
|
||||
// outAlpha, outBeta};
|
||||
|
||||
// dataFormat: 0 = [sL, bS, nIn]
|
||||
// directionMode: 0 = forward, 1 = backward, 2 = bidirectional sum, 3 = bidirectional concat
|
||||
|
||||
const int dataFormat = params[0];
|
||||
const int directionMode = params[1];
|
||||
|
||||
const int sL = x->sizeAt(0); // dataFormat == 0 ? x->sizeAt(0) : x->sizeAt(1);
|
||||
const int bS = x->sizeAt(1); // dataFormat == 0 ? x->sizeAt(1) : x->sizeAt(0);
|
||||
const int nIn = x->sizeAt(-1);
|
||||
const int nOut = Wx->sizeAt(-1);
|
||||
|
||||
const int dirDim = directionMode < 2 ? 1 : 2; // number of dimensionss, 1 unidirectional, 2 for bidirectional
|
||||
const int hDirDim =
|
||||
directionMode <= 2 ? 1 : 2; // for h array, take into account bidirectional_sum mode (directionMode == 2)
|
||||
|
||||
// evaluate direction
|
||||
rnn_direction direction;
|
||||
switch (directionMode) {
|
||||
case 0:
|
||||
direction = rnn_direction::unidirectional_left2right;
|
||||
break;
|
||||
case 1:
|
||||
direction = rnn_direction::unidirectional_right2left;
|
||||
break;
|
||||
case 2:
|
||||
direction = rnn_direction::bidirectional_sum;
|
||||
break;
|
||||
default:
|
||||
direction = rnn_direction::bidirectional_concat;
|
||||
}
|
||||
|
||||
auto engine = onednnUtils::getEngine(LaunchContext::defaultContext()->engine());
|
||||
|
||||
dnnl::memory::desc x_user_md, wx_user_md, wr_user_md, b_user_md, hI_user_md, cI_user_md, h_user_md, hL_user_md,
|
||||
cL_user_md, x_lstm_md, wx_lstm_md, wr_lstm_md, b_lstm_md, hI_lstm_md, cI_lstm_md, h_lstm_md, hL_lstm_md,
|
||||
cL_lstm_md;
|
||||
|
||||
// input type
|
||||
dnnl::memory::data_type xType;
|
||||
if (x->dataType() == DataType::FLOAT32)
|
||||
xType = dnnl::memory::data_type::f32;
|
||||
else if (x->dataType() == DataType::HALF)
|
||||
xType = dnnl::memory::data_type::f16;
|
||||
else
|
||||
xType = dnnl::memory::data_type::u8;
|
||||
|
||||
// weights type
|
||||
dnnl::memory::data_type wType = xType;
|
||||
if (xType == dnnl::memory::data_type::u8) wType = dnnl::memory::data_type::s8;
|
||||
|
||||
// bias type
|
||||
dnnl::memory::data_type bType = xType;
|
||||
if (xType == dnnl::memory::data_type::u8) bType = dnnl::memory::data_type::f32;
|
||||
|
||||
// output type
|
||||
dnnl::memory::data_type hType;
|
||||
if (h->dataType() == DataType::FLOAT32)
|
||||
hType = dnnl::memory::data_type::f32;
|
||||
else if (h->dataType() == DataType::HALF)
|
||||
hType = dnnl::memory::data_type::f16;
|
||||
else
|
||||
hType = dnnl::memory::data_type::u8;
|
||||
|
||||
// memory descriptors for arrays
|
||||
// x
|
||||
x_lstm_md = dnnl::memory::desc({sL, bS, nIn}, xType, dnnl::memory::format_tag::any);
|
||||
// x_user_md = dataFormat == 0 ? dnnl::memory::desc({sL, bS, nIn}, type, dnnl::memory::format_tag::tnc) :
|
||||
// dnnl::memory::desc({bS, sL, nIn}, type, dnnl::memory::format_tag::ntc);
|
||||
x_user_md = dnnl::memory::desc({sL, bS, nIn}, xType, dnnl::memory::format_tag::tnc);
|
||||
onednnUtils::setBlockStrides(*x, x_user_md);
|
||||
|
||||
// wx
|
||||
wx_lstm_md = dnnl::memory::desc({1, dirDim, nIn, 4, nOut}, wType, dnnl::memory::format_tag::any);
|
||||
wx_user_md = dnnl::memory::desc({1, dirDim, nIn, 4, nOut}, wType, dnnl::memory::format_tag::ldigo);
|
||||
onednnUtils::setBlockStrides(*Wx, wx_user_md);
|
||||
|
||||
// wr
|
||||
wr_lstm_md = dnnl::memory::desc({1, dirDim, nOut, 4, nOut}, wType, dnnl::memory::format_tag::any);
|
||||
wr_user_md = dnnl::memory::desc({1, dirDim, nOut, 4, nOut}, wType, dnnl::memory::format_tag::ldigo);
|
||||
onednnUtils::setBlockStrides(*Wr, wr_user_md);
|
||||
|
||||
// h
|
||||
h_lstm_md = dnnl::memory::desc({sL, bS, hDirDim * nOut}, hType, dnnl::memory::format_tag::any);
|
||||
// h_user_md = dataFormat == 0 ? dnnl::memory::desc({sL, bS, hDirDim*nOut}, type, dnnl::memory::format_tag::tnc) :
|
||||
// dnnl::memory::desc({bS, sL, hDirDim*nOut}, type, dnnl::memory::format_tag::ntc);
|
||||
h_user_md = dnnl::memory::desc({sL, bS, hDirDim * nOut}, hType, dnnl::memory::format_tag::tnc);
|
||||
onednnUtils::setBlockStrides(*h, h_user_md);
|
||||
|
||||
// b
|
||||
if (b) {
|
||||
b_lstm_md = dnnl::memory::desc({1, dirDim, 4, nOut}, bType, dnnl::memory::format_tag::any);
|
||||
b_user_md = dnnl::memory::desc({1, dirDim, 4, nOut}, bType, dnnl::memory::format_tag::ldgo);
|
||||
onednnUtils::setBlockStrides(*b, b_user_md);
|
||||
}
|
||||
|
||||
// hI
|
||||
if (hI) {
|
||||
hI_lstm_md = dnnl::memory::desc({1, dirDim, bS, nOut}, xType, dnnl::memory::format_tag::any);
|
||||
hI_user_md = dnnl::memory::desc({1, dirDim, bS, nOut}, xType, dnnl::memory::format_tag::ldnc);
|
||||
onednnUtils::setBlockStrides(*hI, hI_user_md);
|
||||
}
|
||||
|
||||
// cI
|
||||
if (cI) {
|
||||
cI_lstm_md = dnnl::memory::desc({1, dirDim, bS, nOut}, xType, dnnl::memory::format_tag::any);
|
||||
cI_user_md = dnnl::memory::desc({1, dirDim, bS, nOut}, xType, dnnl::memory::format_tag::ldnc);
|
||||
onednnUtils::setBlockStrides(*cI, cI_user_md);
|
||||
}
|
||||
|
||||
// hL
|
||||
if (hL) {
|
||||
hL_lstm_md = dnnl::memory::desc({1, dirDim, bS, nOut}, hType, dnnl::memory::format_tag::any);
|
||||
hL_user_md = dnnl::memory::desc({1, dirDim, bS, nOut}, hType, dnnl::memory::format_tag::ldnc);
|
||||
hL_user_md.data.format_kind = dnnl_blocked; // overrides format
|
||||
onednnUtils::setBlockStrides(*hL, hL_user_md);
|
||||
}
|
||||
|
||||
if (cL) {
|
||||
cL_lstm_md = dnnl::memory::desc({1, dirDim, bS, nOut}, hType, dnnl::memory::format_tag::ldnc);
|
||||
cL_user_md = dnnl::memory::desc({1, dirDim, bS, nOut}, hType, dnnl::memory::format_tag::ldnc);
|
||||
onednnUtils::setBlockStrides(*cL, cL_user_md);
|
||||
}
|
||||
|
||||
// lstm memory description
|
||||
lstm_forward::desc lstm_desc(prop_kind::forward_inference, direction, x_lstm_md, hI_lstm_md, cI_lstm_md, wx_lstm_md,
|
||||
wr_lstm_md, b_lstm_md, h_lstm_md, hL_lstm_md, cL_lstm_md);
|
||||
|
||||
dnnl::stream stream(engine);
|
||||
|
||||
// lstm primitive description
|
||||
lstm_forward::primitive_desc lstm_prim_desc(lstm_desc, engine);
|
||||
|
||||
// arguments (memory buffers) necessary for calculations
|
||||
std::unordered_map<int, dnnl::memory> args;
|
||||
|
||||
// provide memory and check whether reorder is required
|
||||
// x
|
||||
onednnUtils::loadDataToMklStream(*x, engine, stream, x_user_md, lstm_prim_desc.src_layer_desc(),
|
||||
args[DNNL_ARG_SRC_LAYER]);
|
||||
|
||||
// wx
|
||||
onednnUtils::loadDataToMklStream(*Wx, engine, stream, wx_user_md, lstm_prim_desc.weights_layer_desc(),
|
||||
args[DNNL_ARG_WEIGHTS_LAYER]);
|
||||
|
||||
// wr
|
||||
onednnUtils::loadDataToMklStream(*Wr, engine, stream, wr_user_md, lstm_prim_desc.weights_iter_desc(),
|
||||
args[DNNL_ARG_WEIGHTS_ITER]);
|
||||
|
||||
// h
|
||||
auto h_user_mem = onednnUtils::loadDataToMklStream(*h, engine, stream, h_user_md, lstm_prim_desc.dst_layer_desc(),
|
||||
args[DNNL_ARG_DST_LAYER]);
|
||||
|
||||
// b
|
||||
if (b)
|
||||
onednnUtils::loadDataToMklStream(*b, engine, stream, b_user_md, lstm_prim_desc.bias_desc(), args[DNNL_ARG_BIAS]);
|
||||
|
||||
// hI
|
||||
if (hI)
|
||||
onednnUtils::loadDataToMklStream(*hI, engine, stream, hI_user_md, lstm_prim_desc.src_iter_desc(),
|
||||
args[DNNL_ARG_SRC_ITER]);
|
||||
|
||||
// cI
|
||||
if (cI)
|
||||
onednnUtils::loadDataToMklStream(*cI, engine, stream, cI_user_md, lstm_prim_desc.src_iter_c_desc(),
|
||||
args[DNNL_ARG_SRC_ITER_C]);
|
||||
|
||||
dnnl::memory hL_user_mem, cL_user_mem, hL_lstm_mem, cL_lstm_mem;
|
||||
|
||||
// hL
|
||||
if (hL)
|
||||
hL_user_mem = onednnUtils::loadDataToMklStream(*hL, engine, stream, hL_user_md, lstm_prim_desc.dst_iter_desc(),
|
||||
args[DNNL_ARG_DST_ITER]);
|
||||
|
||||
// cL
|
||||
if (cL)
|
||||
cL_user_mem = onednnUtils::loadDataToMklStream(*cL, engine, stream, cL_user_md, lstm_prim_desc.dst_iter_c_desc(),
|
||||
args[DNNL_ARG_DST_ITER_C]);
|
||||
|
||||
// run calculations
|
||||
lstm_forward(lstm_prim_desc).execute(stream, args);
|
||||
|
||||
// reorder outputs if necessary
|
||||
if (lstm_prim_desc.dst_layer_desc() != h_user_mem.get_desc())
|
||||
reorder(args[DNNL_ARG_DST_LAYER], h_user_mem).execute(stream, args[DNNL_ARG_DST_LAYER], h_user_mem);
|
||||
if (lstm_prim_desc.dst_iter_desc() != hL_user_mem.get_desc())
|
||||
reorder(args[DNNL_ARG_DST_ITER], hL_user_mem).execute(stream, args[DNNL_ARG_DST_ITER], hL_user_mem);
|
||||
if (lstm_prim_desc.dst_iter_c_desc() != cL_user_mem.get_desc())
|
||||
reorder(args[DNNL_ARG_DST_ITER_C], cL_user_mem).execute(stream, args[DNNL_ARG_DST_ITER_C], cL_user_mem);
|
||||
|
||||
stream.wait();
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
PLATFORM_IMPL(lstmLayer, ENGINE_CPU) {
|
||||
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)
|
||||
|
||||
const auto hasBiases = B_ARG(0); // indicates whether biases array is provided
|
||||
const auto hasSeqLen = 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 seqLen = hasSeqLen ? 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
|
||||
|
||||
REQUIRE_TRUE(cellClip == 0, 0, "LSTM_LAYER_MKLDNN operation: cell clipping is not supported currently !");
|
||||
REQUIRE_TRUE(retFullSeq, 0,
|
||||
"LSTM_LAYER_MKLDNN operation: option to calculate full time sequence output h should be always true in "
|
||||
"case of mkl dnn library !");
|
||||
REQUIRE_TRUE(hasPH == false, 0,
|
||||
"LSTM_LAYER_MKLDNN operation: mkl dnn library doesn't support peephole connections !");
|
||||
REQUIRE_TRUE(hasSeqLen == false, 0,
|
||||
"LSTM_LAYER_MKLDNN operation: mkl dnn library doesn't support array specifying max time step per each "
|
||||
"example in batch !");
|
||||
REQUIRE_TRUE(dataFormat < 2, 0,
|
||||
"LSTM_LAYER_MKLDNN operation: wrong data format, only two formats are allowed for input/output tensors "
|
||||
"in mkl dnn library: TNC and NTC!");
|
||||
REQUIRE_TRUE(
|
||||
directionMode < 4, 0,
|
||||
"LSTM_LAYER_MKLDNN operation: option for bidirectional extra output dimension is not valid in mkl dnn library !");
|
||||
REQUIRE_TRUE(retLastH == retLastC, 0,
|
||||
"LSTM_LAYER_MKLDNN operation: only two options are present: 1) calculate both output at last time and "
|
||||
"cell state at last time; 2) do not calculate both !");
|
||||
REQUIRE_TRUE(hasInitH == hasInitC, 0,
|
||||
"LSTM_LAYER_MKLDNN operation: either both of or neither of initial C and initial H must be provided");
|
||||
|
||||
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
|
||||
|
||||
// evaluate dimensions
|
||||
const sd::LongType sL = x->sizeAt(dataFormat);
|
||||
const sd::LongType bS = dataFormat == 0 ? x->sizeAt(1) : x->sizeAt(0);
|
||||
const sd::LongType nIn = x->sizeAt(2);
|
||||
const sd::LongType nOut = Wx->sizeAt(-1) / 4;
|
||||
|
||||
// inputs validations
|
||||
if (directionMode < 2) { // no bidirectional
|
||||
|
||||
// Wx validation
|
||||
if (Wx->rankOf() != 2 || Wx->sizeAt(0) != nIn)
|
||||
REQUIRE_TRUE(false, 0,
|
||||
"LSTM_LAYER_MKLDNN 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_MKLDNN 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_MKLDNN 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_MKLDNN 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_MKLDNN 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_MKLDNN 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_MKLDNN 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_MKLDNN 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_MKLDNN 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_MKLDNN 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());
|
||||
}
|
||||
|
||||
std::vector<float> params = {static_cast<float>(dataFormat), static_cast<float>(directionMode),
|
||||
static_cast<float>(cellClip)};
|
||||
|
||||
const int dirDim = directionMode < 2 ? 1 : 2; // number of dimensions, 1 unidirectional, 2 for bidirectional
|
||||
|
||||
// permut x and h to tnc format if they have ntc format
|
||||
NDArray *xP(const_cast<NDArray*>(x)), *hP(h);
|
||||
if (dataFormat == 1) {
|
||||
std::vector<sd::LongType> permute = {1,0,2};
|
||||
xP = new NDArray(x->permute(permute.data(), 3, false, false)); // [bS, sL, nIn] -> [sL, bS, nIn]
|
||||
hP = new NDArray(h->permute(permute.data(), 3, false, false)); // [bS, sL, dirDim*nOn] -> [sL, bS, dirDim*nOn]
|
||||
}
|
||||
|
||||
// reshape arrays in accordance to mkl allowed formats
|
||||
NDArray *WxR(nullptr), *WrR(nullptr), *bR(nullptr), *hIR(nullptr), *cIR(nullptr), *hLR(nullptr), *cLR(nullptr);
|
||||
|
||||
std::vector<sd::LongType> shapeOne = {1, dirDim, nIn, 4, nOut};
|
||||
WxR = new NDArray(Wx->reshape(Wx->ordering(), shapeOne));
|
||||
WrR = new NDArray(Wr->reshape(Wr->ordering(), shapeOne));
|
||||
|
||||
std::vector<sd::LongType> shapeTwo = {1, dirDim, 4, nOut};
|
||||
if (b)
|
||||
bR = new NDArray(b->reshape(b->ordering(), shapeTwo));
|
||||
else
|
||||
bR =
|
||||
new NDArray(x->ordering(), shapeTwo, x->dataType(), x->getContext()); // already nullified
|
||||
|
||||
|
||||
std::vector<sd::LongType> shapeThree = {1, dirDim, bS, nOut};
|
||||
if (hI) hIR = new NDArray(hI->reshape(hI->ordering(), shapeThree));
|
||||
|
||||
if (cI) cIR = new NDArray(cI->reshape(cI->ordering(), shapeThree));
|
||||
|
||||
if (hL) hLR = new NDArray(hL->reshape(hL->ordering(), shapeThree, false));
|
||||
|
||||
if (cL) cLR = new NDArray(cL->reshape(cL->ordering(), shapeThree, false));
|
||||
|
||||
lstmLayerMKLDNN(xP, WxR, WrR, bR, hIR, cIR, params, hP, hLR, cLR);
|
||||
|
||||
delete WxR;
|
||||
delete WrR;
|
||||
delete bR;
|
||||
delete hIR;
|
||||
delete cIR;
|
||||
delete hLR;
|
||||
delete cLR;
|
||||
|
||||
if (dataFormat == 1) {
|
||||
delete xP;
|
||||
delete hP;
|
||||
}
|
||||
|
||||
return sd::Status::OK;
|
||||
}
|
||||
|
||||
PLATFORM_CHECK(lstmLayer, ENGINE_CPU) {
|
||||
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)
|
||||
|
||||
const auto hasBiases = B_ARG(0); // indicates whether biases array is provided
|
||||
const auto hasSeqLen = 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();
|
||||
DataType bType = b != nullptr ? b->dataType() : (xType == DataType::HALF ? xType : DataType::FLOAT32);
|
||||
DataType hIType = hI != nullptr ? hI->dataType() : xType;
|
||||
DataType cIType = cI != nullptr ? cI->dataType() : xType;
|
||||
DataType hType = h != nullptr ? h->dataType() : xType;
|
||||
DataType hLType = hL != nullptr ? hL->dataType() : xType;
|
||||
DataType cLType = cL != nullptr ? cL->dataType() : xType;
|
||||
Requirements req("ONEDNN LstmLayer OP");
|
||||
req.expectTrue(block.isUseONEDNN(), IS_USE_ONEDNN_MSG) &&
|
||||
req.expectEq(makeInfoVariable(cellClip, MSG_CELL_CLIPPING), 0) &&
|
||||
req.expectTrue(makeInfoVariable(retFullSeq, "Return full sequence")) &&
|
||||
req.expectFalse(makeInfoVariable(hasPH, HAVE_PEEPHOLE), EXPECTED_NOT_SUPPORTED) &&
|
||||
req.expectFalse(makeInfoVariable(hasSeqLen, HAVE_SEQLENARR), EXPECTED_NOT_SUPPORTED) &&
|
||||
req.expectLess(makeInfoVariable(dataFormat, "Data format"), 2) &&
|
||||
req.expectLess(makeInfoVariable(directionMode, "Direction mode"), 4) &&
|
||||
req.expectEq(makeInfoVariable(retLastH, "Return lastH"), makeInfoVariable(retLastC, "Return lastC")) &&
|
||||
req.expectEq(makeInfoVariable(hasInitH, "Has initial H"), makeInfoVariable(hasInitC, "Has initial C")) &&
|
||||
req.expectTrue(
|
||||
makeInfoVariable(
|
||||
[xType, WxType, WrType, bType, hIType, cIType, hType, hLType, cLType] {
|
||||
return ((xType == DataType::FLOAT32 && WxType == DataType::FLOAT32 && WrType == DataType::FLOAT32 &&
|
||||
bType == DataType::FLOAT32 && hIType == DataType::FLOAT32 && cIType == DataType::FLOAT32 &&
|
||||
hType == DataType::FLOAT32 && hLType == DataType::FLOAT32 && cLType == DataType::FLOAT32) ||
|
||||
(xType == DataType::HALF && WxType == DataType::HALF && WrType == DataType::HALF &&
|
||||
bType == DataType::HALF && hIType == DataType::HALF && cIType == DataType::HALF &&
|
||||
hType == DataType::HALF && hLType == DataType::HALF && cLType == DataType::HALF) ||
|
||||
(xType == DataType::UINT8 && WxType == DataType::INT8 && WrType == DataType::INT8 &&
|
||||
bType == DataType::FLOAT32 && hIType == DataType::UINT8 && cIType == DataType::UINT8 &&
|
||||
((hType == DataType::FLOAT32 && hLType == DataType::FLOAT32 && cLType == DataType::FLOAT32) ||
|
||||
(hType == DataType::UINT8 && hLType == DataType::UINT8 && cLType == DataType::UINT8))));
|
||||
},
|
||||
TYPECHECK_MSG),
|
||||
NO_MSG);
|
||||
|
||||
req.logTheSuccess();
|
||||
return req;
|
||||
}
|
||||
|
||||
} // namespace platforms
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,328 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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 <helpers/MKLDNNStream.h>
|
||||
#include <ops/declarable/OpRegistrator.h>
|
||||
#include <ops/declarable/PlatformHelper.h>
|
||||
#include <system/platform_boilerplate.h>
|
||||
|
||||
#include <numeric>
|
||||
|
||||
#include "mkldnnUtils.h"
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace platforms {
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
static void matmulMKLDNN(NDArray* x, NDArray* y, NDArray* z, const bool transX, const bool transY,
|
||||
float alpha = 1.f, float beta = 0.f) {
|
||||
// mkl works with following
|
||||
// [M,K] x [K,N] = [M,N]
|
||||
// [bS, M,K] x [bS, K,N] = [bS, M,N]
|
||||
|
||||
// possible input cases not supported by mkl, however we'll perform permut/reshape procedures in order to fit
|
||||
// requirements [4] x [4] = [1] --> [1,4] x [4,1] = [1,1] [4] x [4,5] =
|
||||
// [5] --> [1,4] x [4,5] = [1,5] [4,5] x [5] = [4] --> [4,5] x [5,1] =
|
||||
// [4,1] [2,3, 4,5] x [2,3, 5,4] = [2,3, 4,4] --> [6, 4,5] x [6, 5,4] = [6, 4,4] [2,2,3, 4,5] x [2,2,3, 5,4] =
|
||||
// [2,2,3, 4,4] --> [12, 4,5] x [12, 5,4] = [12, 4,4]
|
||||
|
||||
const auto xRank = x->rankOf();
|
||||
const auto yRank = y->rankOf();
|
||||
const auto zRank = z->rankOf();
|
||||
|
||||
std::vector<sd::LongType> permut;
|
||||
|
||||
// fill permutation vector appropriately if transposition is required
|
||||
if ((transX && xRank > 1) || (transY && yRank > 1)) {
|
||||
const int rank = xRank >= yRank ? xRank : yRank;
|
||||
permut.resize(rank);
|
||||
std::iota(std::begin(permut), std::end(permut), 0);
|
||||
permut[rank - 2] = rank - 1;
|
||||
permut[rank - 1] = rank - 2;
|
||||
}
|
||||
|
||||
NDArray* xT = (transX && xRank > 1) ? new NDArray(x->permute(permut, false, false)) : x;
|
||||
NDArray* yT = (transY && yRank > 1) ? new NDArray(y->permute(permut, false, false)) : y;
|
||||
|
||||
|
||||
std::vector<sd::LongType> shapeOne = {xT->lengthOf() / (xT->sizeAt(-2) * xT->sizeAt(-1)),
|
||||
xT->sizeAt(-2), xT->sizeAt(-1)};
|
||||
NDArray* xTR =
|
||||
xRank <= 3 ? xT
|
||||
: new NDArray(xT->reshape(xT->ordering(),shapeOne));
|
||||
std::vector<sd::LongType> shapeTwo = {yT->lengthOf() / (yT->sizeAt(-2) * yT->sizeAt(-1)),
|
||||
yT->sizeAt(-2), yT->sizeAt(-1)};
|
||||
NDArray* yTR =
|
||||
xRank <= 3 ? yT
|
||||
: new NDArray(yT->reshape(yT->ordering(),shapeTwo));
|
||||
std::vector<sd::LongType> shapeThree = {z->lengthOf() / (z->sizeAt(-2) * z->sizeAt(-1)),
|
||||
z->sizeAt(-2), z->sizeAt(-1)};
|
||||
NDArray* zR = xRank <= 3 ? z
|
||||
: new NDArray(z->reshape(z->ordering(), shapeThree) /*, false*/);
|
||||
|
||||
// [M,K] x [K,N] = [M,N]
|
||||
const sd::LongType M = (xRank > 1) ? xTR->sizeAt(-2) : 1;
|
||||
const sd::LongType K = (xRank > 1) ? xTR->sizeAt(-1) : xTR->lengthOf();
|
||||
const sd::LongType N = (yRank > 1) ? yTR->sizeAt(-1) : 1;
|
||||
const sd::LongType bS = (xRank > 2) ? xTR->sizeAt(0) : 1; // [bS, M,K] x [bS, K,N] = [bS, M,N]
|
||||
|
||||
dnnl::memory::dims xShape = xRank < 3 ? dnnl::memory::dims({M, K}) : dnnl::memory::dims({bS, M, K});
|
||||
dnnl::memory::dims yShape = xRank < 3 ? dnnl::memory::dims({K, N}) : dnnl::memory::dims({bS, K, N});
|
||||
dnnl::memory::dims zShape = xRank < 3 ? dnnl::memory::dims({M, N}) : dnnl::memory::dims({bS, M, N});
|
||||
|
||||
// x type
|
||||
dnnl::memory::data_type xType;
|
||||
if (x->dataType() == DataType::FLOAT32)
|
||||
xType = dnnl::memory::data_type::f32;
|
||||
else if (x->dataType() == DataType::HALF)
|
||||
xType = dnnl::memory::data_type::f16;
|
||||
else if (x->dataType() == DataType::BFLOAT16)
|
||||
xType = dnnl::memory::data_type::bf16;
|
||||
else if (x->dataType() == DataType::UINT8)
|
||||
xType = dnnl::memory::data_type::u8;
|
||||
else
|
||||
xType = dnnl::memory::data_type::s8;
|
||||
|
||||
// y type
|
||||
dnnl::memory::data_type yType = xType;
|
||||
if (y->dataType() == DataType::UINT8)
|
||||
yType = dnnl::memory::data_type::u8;
|
||||
else if (y->dataType() == DataType::INT8)
|
||||
yType = dnnl::memory::data_type::s8;
|
||||
|
||||
// z type
|
||||
dnnl::memory::data_type zType = xType;
|
||||
if (z->dataType() == DataType::FLOAT32)
|
||||
zType = dnnl::memory::data_type::f32;
|
||||
else if (z->dataType() == DataType::INT32)
|
||||
zType = dnnl::memory::data_type::s32;
|
||||
else if (z->dataType() == DataType::UINT8)
|
||||
zType = dnnl::memory::data_type::u8;
|
||||
else if (z->dataType() == DataType::INT8)
|
||||
zType = dnnl::memory::data_type::s8;
|
||||
|
||||
const auto xFormat = xRank == 1 ? dnnl::memory::format_tag::ab : onednnUtils::getFormat(*xTR);
|
||||
const auto yFormat = yRank == 1 ? dnnl::memory::format_tag::ab : onednnUtils::getFormat(*yTR);
|
||||
const auto zFormat = zRank == 1 ? dnnl::memory::format_tag::ab : onednnUtils::getFormat(*zR);
|
||||
|
||||
// memory descriptors for arrays
|
||||
dnnl::memory::desc x_mkl_md, x_user_md, y_mkl_md, y_user_md, z_mkl_md, z_user_md;
|
||||
|
||||
// x
|
||||
x_user_md = x_mkl_md = dnnl::memory::desc(xShape, xType, xFormat);
|
||||
x_user_md.data.format_kind = dnnl_blocked; // overrides format
|
||||
x_user_md.data.format_desc.blocking.strides[0] = xRank == 1 ? 1 : xTR->strideAt(0);
|
||||
x_user_md.data.format_desc.blocking.strides[1] = xRank == 1 ? xTR->strideAt(0) : xTR->strideAt(1);
|
||||
if (xRank > 2) x_user_md.data.format_desc.blocking.strides[2] = xTR->strideAt(2);
|
||||
|
||||
|
||||
// y
|
||||
y_user_md = y_mkl_md = dnnl::memory::desc(yShape, yType, yFormat);
|
||||
y_user_md.data.format_kind = dnnl_blocked; // overrides format
|
||||
y_user_md.data.format_desc.blocking.strides[0] = yRank == 1 ? 1 : yTR->strideAt(0);
|
||||
y_user_md.data.format_desc.blocking.strides[1] = yRank == 1 ? yTR->strideAt(0) : yTR->strideAt(1);
|
||||
if (yRank > 2) y_user_md.data.format_desc.blocking.strides[2] = yTR->strideAt(2);
|
||||
|
||||
|
||||
// z
|
||||
z_user_md = z_mkl_md = dnnl::memory::desc(zShape, zType, zFormat);
|
||||
z_user_md.data.format_kind = dnnl_blocked; // overrides format
|
||||
z_user_md.data.format_desc.blocking.strides[0] = zRank == 1 ? 1 : zR->strideAt(0);
|
||||
z_user_md.data.format_desc.blocking.strides[1] = zRank == 1 ? zR->strideAt(0) : zR->strideAt(1);
|
||||
if (zRank > 2) z_user_md.data.format_desc.blocking.strides[2] = zR->strideAt(2);
|
||||
|
||||
|
||||
auto engine = onednnUtils::getEngine(LaunchContext::defaultContext()->engine());
|
||||
|
||||
// Create attributes (to handle alpha and beta if necessary)
|
||||
dnnl::primitive_attr attr; // it is empty since we have usual values for alpha (=1) and beta (=0)
|
||||
if (alpha != 1.f) attr.set_output_scales(0, {alpha});
|
||||
if (beta != 0.f) {
|
||||
dnnl::post_ops po;
|
||||
po.append_sum(beta);
|
||||
attr.set_post_ops(po);
|
||||
}
|
||||
|
||||
// operation primitive description
|
||||
dnnl::matmul::desc op_desc(x_mkl_md, y_mkl_md, z_mkl_md);
|
||||
dnnl::matmul::primitive_desc op_prim_desc(op_desc, attr, engine);
|
||||
|
||||
// arguments (memory buffers) necessary for calculations
|
||||
std::unordered_map<int, dnnl::memory> args;
|
||||
|
||||
dnnl::stream stream(engine);
|
||||
|
||||
// provide memory buffers and check whether reorder is required
|
||||
|
||||
// input
|
||||
onednnUtils::loadDataToMklStream(*xTR, engine, stream, x_user_md, op_prim_desc.src_desc(), args[DNNL_ARG_SRC]);
|
||||
|
||||
// y
|
||||
onednnUtils::loadDataToMklStream(*yTR, engine, stream, y_user_md, op_prim_desc.weights_desc(),
|
||||
args[DNNL_ARG_WEIGHTS]);
|
||||
|
||||
// z
|
||||
auto z_user_mem =
|
||||
onednnUtils::loadDataToMklStream(*zR, engine, stream, z_user_md, op_prim_desc.dst_desc(), args[DNNL_ARG_DST]);
|
||||
|
||||
// run calculations
|
||||
dnnl::matmul(op_prim_desc).execute(stream, args);
|
||||
|
||||
// reorder outputs if necessary
|
||||
if (op_prim_desc.dst_desc() != z_user_mem.get_desc())
|
||||
dnnl::reorder(args[DNNL_ARG_DST], z_user_mem).execute(stream, args[DNNL_ARG_DST], z_user_mem);
|
||||
|
||||
stream.wait();
|
||||
|
||||
if (zR->buffer() != z->buffer()) z->assign(zR);
|
||||
|
||||
if (zR != z) delete zR;
|
||||
if (xTR != xT) delete xTR;
|
||||
if (xT != x) delete xT;
|
||||
if (yTR != yT) delete yTR;
|
||||
if (yT != y) delete yT;
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
PLATFORM_IMPL(matmul, ENGINE_CPU) {
|
||||
auto x = INPUT_VARIABLE(0);
|
||||
auto y = INPUT_VARIABLE(1);
|
||||
auto z = OUTPUT_VARIABLE(0);
|
||||
|
||||
if (x->isEmpty() || y->isEmpty()) return sd::Status::OK;
|
||||
|
||||
sd::LongType iSize = (sd::LongType)block.getIArguments()->size();
|
||||
int transX = iSize > 0 ? INT_ARG(0) : 0;
|
||||
int transY = iSize > 1 ? INT_ARG(1) : 0;
|
||||
const int transZ = iSize > 2 ? INT_ARG(2) : 0;
|
||||
|
||||
// optional use alpha nad beta
|
||||
iSize = (sd::LongType)block.getTArguments()->size();
|
||||
float alpha = iSize > 0 ? T_ARG(0) : 1.0;
|
||||
float beta = iSize > 1 ? T_ARG(1) : 0.0;
|
||||
|
||||
const sd::LongType xRank = x->rankOf();
|
||||
const sd::LongType yRank = y->rankOf();
|
||||
const sd::LongType zRank = z->rankOf();
|
||||
|
||||
if (transZ) {
|
||||
x = INPUT_VARIABLE(1);
|
||||
y = INPUT_VARIABLE(0);
|
||||
bool temp = transX;
|
||||
transX = !transY;
|
||||
transY = !temp;
|
||||
}
|
||||
|
||||
const sd::LongType xLastDim = transX ? -2 : -1;
|
||||
const sd::LongType yLastDim = transY ? -2 : -1;
|
||||
const sd::LongType xLastButOneDim = transX ? -1 : -2;
|
||||
const sd::LongType yLastButOneDim = transY ? -1 : -2;
|
||||
|
||||
// ******* input validation ******* //
|
||||
REQUIRE_TRUE(xRank > 0 && yRank > 0, 0,
|
||||
"MATMUL MKLDNN OP: input arrays must have rank bigger than 0 (should not be scalars), but got instead: "
|
||||
"x rank = %i, y rank = %i !",
|
||||
xRank, yRank);
|
||||
|
||||
if (xRank == 1 && yRank == 1) { // dot case, output is scalar (or vector with length = 1)
|
||||
REQUIRE_TRUE(x->lengthOf() == y->lengthOf(), 0,
|
||||
"MATMUL MKLDNN OP: since input arrays are vectors they must have the same length, but got x length = "
|
||||
"%i, y length = %i !",
|
||||
x->lengthOf(), y->lengthOf());
|
||||
} else if (xRank == 1 && yRank == 2) { // vector x matrix, i.e. [4] x [4,5] = [5], output is vector
|
||||
REQUIRE_TRUE(x->lengthOf() == y->sizeAt(yLastButOneDim), 0,
|
||||
"MATMUL MKLDNN OP: input arrays have inconsistent shapes for vector-matrix product: x %s, y %s !",
|
||||
ShapeUtils::shapeAsString(x).c_str(), ShapeUtils::shapeAsString(y).c_str());
|
||||
} else if (xRank == 2 && yRank == 1) { // matrix x vector , i.e. [4,5] x [5] = [4], output is vector
|
||||
REQUIRE_TRUE(x->sizeAt(xLastDim) == y->lengthOf(), 0,
|
||||
"MATMUL MKLDNN OP: input arrays have inconsistent shapes for matrix-vector product: x %s, y %s !",
|
||||
ShapeUtils::shapeAsString(x).c_str(), ShapeUtils::shapeAsString(y).c_str());
|
||||
} else {
|
||||
REQUIRE_TRUE(xRank == yRank && yRank == zRank, 0,
|
||||
"MATMUL MKLDNN OP: input and output arrays must have the same rank, but got instead: x rank = %i, y "
|
||||
"rank = %i, z rank = %i !",
|
||||
xRank, yRank, zRank);
|
||||
REQUIRE_TRUE(
|
||||
x->sizeAt(xLastDim) == y->sizeAt(yLastButOneDim) && x->sizeAt(xLastButOneDim) == z->sizeAt(-2) &&
|
||||
y->sizeAt(yLastDim) == z->sizeAt(-1),
|
||||
0, "MATMUL MKLDNN OP: input/output arrays have inconsistent shapes for matrix product: x %s, y %s, z %s !",
|
||||
ShapeUtils::shapeAsString(x).c_str(), ShapeUtils::shapeAsString(y).c_str(),
|
||||
ShapeUtils::shapeAsString(z).c_str());
|
||||
|
||||
if (xRank > 2) // outer dims must be the same
|
||||
for (int i = 0; i < xRank - 2; ++i)
|
||||
REQUIRE_TRUE(
|
||||
x->sizeAt(i) == y->sizeAt(i) && y->sizeAt(i) == z->sizeAt(i), 0,
|
||||
"MATMUL MKLDNN OP: input/output arrays have inconsistent shapes for matrix product: x %s, y %s, z %s !",
|
||||
ShapeUtils::shapeAsString(x).c_str(), ShapeUtils::shapeAsString(y).c_str(),
|
||||
ShapeUtils::shapeAsString(z).c_str());
|
||||
}
|
||||
// ******* end of input validation ******* //
|
||||
|
||||
matmulMKLDNN(x, y, z, transX, transY, alpha, beta);
|
||||
|
||||
return sd::Status::OK;
|
||||
}
|
||||
#include <iostream>
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
PLATFORM_CHECK(matmul, ENGINE_CPU) {
|
||||
auto x = INPUT_VARIABLE(0);
|
||||
auto y = INPUT_VARIABLE(1);
|
||||
|
||||
auto z = OUTPUT_VARIABLE(0);
|
||||
|
||||
const auto xType = x->dataType();
|
||||
const auto yType = y->dataType();
|
||||
const auto zType = z->dataType();
|
||||
|
||||
float alpha = block.numT() > 0 ? T_ARG(0) : 1.0f;
|
||||
float beta = block.numT() > 1 ? T_ARG(1) : 0.0f;
|
||||
|
||||
Requirements req("ONEDNN MATMUL OP");
|
||||
|
||||
// we're skipping if result order is F or arrays are not continuous
|
||||
req.expectTrue(block.isUseONEDNN(), IS_USE_ONEDNN_MSG) &&
|
||||
req.expectLess(makeInfoVariable(x->rankOf(), RANK_MSG_INPUT0), 3);
|
||||
|
||||
req.setPrefix("ONEDNN MATMUL OP")
|
||||
.expectTrue(
|
||||
makeInfoVariable(
|
||||
[xType, yType, zType] {
|
||||
return ((xType == DataType::FLOAT32 && yType == DataType::FLOAT32 && zType == DataType::FLOAT32) ||
|
||||
(xType == DataType::HALF && yType == DataType::HALF && zType == DataType::FLOAT32) ||
|
||||
(xType == DataType::BFLOAT16 && yType == DataType::BFLOAT16 && zType == DataType::BFLOAT16) ||
|
||||
((xType == DataType::UINT8 || xType == DataType::INT8) &&
|
||||
(yType == DataType::UINT8 || yType == DataType::INT8) &&
|
||||
(zType == DataType::UINT8 || zType == DataType::INT8 || zType == DataType::INT32 ||
|
||||
zType == DataType::FLOAT32)));
|
||||
},
|
||||
TYPECHECK_MSG),
|
||||
NO_MSG);
|
||||
|
||||
req.logTheSuccess();
|
||||
|
||||
return req;
|
||||
}
|
||||
|
||||
} // namespace platforms
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,148 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* This program and the accompanying materials are made available under the
|
||||
* terms of the Apache License, Version 2.0 which is available at
|
||||
* https://www.apache.org/licenses/LICENSE-2.0.
|
||||
*
|
||||
* See the NOTICE file distributed with this work for additional
|
||||
* information regarding copyright ownership.
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
|
||||
* WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
|
||||
* License for the specific language governing permissions and limitations
|
||||
* under the License.
|
||||
*
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
******************************************************************************/
|
||||
|
||||
//
|
||||
// @author saudet
|
||||
// @author raver119@gmail.com
|
||||
// @author Yurii Shyrma (iuriish@yahoo.com)
|
||||
//
|
||||
#include <helpers/MKLDNNStream.h>
|
||||
#include <ops/declarable/OpRegistrator.h>
|
||||
#include <ops/declarable/PlatformHelper.h>
|
||||
#include <ops/declarable/helpers/convolutions.h>
|
||||
#include <system/platform_boilerplate.h>
|
||||
|
||||
#include "mkldnnUtils.h"
|
||||
|
||||
using namespace dnnl;
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace platforms {
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
PLATFORM_IMPL(maxpool2d, ENGINE_CPU) {
|
||||
auto input = INPUT_VARIABLE(0);
|
||||
auto output = OUTPUT_VARIABLE(0);
|
||||
|
||||
REQUIRE_TRUE(input->rankOf() == 4, 0, "MAXPOOL2D MKLDNN OP: input array should have rank of 4, but got %i instead",
|
||||
input->rankOf());
|
||||
|
||||
// 0,1 - kernel Height/Width; 2,3 - stride Height/Width; 4,5 - pad Height/Width; 6,7 - dilation Height/Width; 8 - same
|
||||
// mode;
|
||||
const sd::LongType kH = INT_ARG(0);
|
||||
const sd::LongType kW = INT_ARG(1);
|
||||
const sd::LongType sH = INT_ARG(2);
|
||||
const sd::LongType sW = INT_ARG(3);
|
||||
sd::LongType pH = INT_ARG(4);
|
||||
sd::LongType pW = INT_ARG(5);
|
||||
const sd::LongType dH = INT_ARG(6);
|
||||
const sd::LongType dW = INT_ARG(7);
|
||||
const int paddingMode = INT_ARG(8);
|
||||
// const int extraParam0 = INT_ARG(9);
|
||||
const int isNCHW = block.getIArguments()->size() > 10 ? !INT_ARG(10) : 1; // INT_ARG(10): 1-NHWC, 0-NCHW
|
||||
|
||||
REQUIRE_TRUE(dH != 0 && dW != 0, 0, "MAXPOOL2D MKLDNN op: dilation must not be zero, but got instead {%i, %i}", dH,
|
||||
dW);
|
||||
|
||||
sd::LongType bS, iC, iH, iW, oC, oH,
|
||||
oW; // batch size, input channels, input height/width, output channels, output height/width;
|
||||
sd::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);
|
||||
|
||||
if (paddingMode) ConvolutionUtils::calcPadding2D(pH, pW, oH, oW, iH, iW, kH, kW, sH, sW, dH, dW);
|
||||
|
||||
onednnUtils::poolingONEDNN(input, output, 0, kH, kW, 0, sH, sW, 0, pH, pW, isNCHW, algorithm::pooling_max);
|
||||
|
||||
return sd::Status::OK;
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
PLATFORM_CHECK(maxpool2d, ENGINE_CPU) {
|
||||
auto input = INPUT_VARIABLE(0);
|
||||
auto output = OUTPUT_VARIABLE(0);
|
||||
|
||||
Requirements req("ONEDNN MAXPOOL2d OP");
|
||||
req.expectTrue(block.isUseONEDNN(), IS_USE_ONEDNN_MSG) &&
|
||||
req.expectTrue(sd::ONEDNNStream::isSupported({input, output}), ONEDNN_STREAM_NOT_SUPPORTED);
|
||||
if (req) onednnUtils::checkPoolingONEDNN(req, block, input, output);
|
||||
req.logTheSuccess();
|
||||
return req;
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
PLATFORM_IMPL(maxpool2d_bp, ENGINE_CPU) {
|
||||
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
|
||||
|
||||
sd::LongType kH = INT_ARG(0); // filter(kernel) height
|
||||
sd::LongType kW = INT_ARG(1); // filter(kernel) width
|
||||
sd::LongType sH = INT_ARG(2); // strides height
|
||||
sd::LongType sW = INT_ARG(3); // strides width
|
||||
sd::LongType pH = INT_ARG(4); // paddings height
|
||||
sd::LongType pW = INT_ARG(5); // paddings width
|
||||
sd::LongType dH = INT_ARG(6); // dilations height
|
||||
sd::LongType dW = INT_ARG(7); // dilations width
|
||||
int paddingMode = INT_ARG(8); // 0-VALID, 1-SAME
|
||||
// int extraParam0 = INT_ARG(9);
|
||||
int isNCHW = block.getIArguments()->size() > 10 ? !INT_ARG(10) : 1; // INT_ARG(10): 0-NCHW, 1-NHWC
|
||||
|
||||
REQUIRE_TRUE(input->rankOf() == 4, 0, "MAXPOOL2D_BP MKLDNN op: input should have rank of 4, but got %i instead",
|
||||
input->rankOf());
|
||||
REQUIRE_TRUE(dH != 0 && dW != 0, 0, "MAXPOOL2D_BP MKLDNN op: dilation must not be zero, but got instead {%i, %i}", dH,
|
||||
dW);
|
||||
|
||||
sd::LongType bS, iC, iH, iW, oC, oH,
|
||||
oW; // batch size, input channels, input height/width, output channels, output height/width;
|
||||
sd::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<sd::LongType> expectedGradOShape =
|
||||
ShapeUtils::composeShapeUsingDimsAndIdx({bS, iC, oH, oW, 0, indIOioC, indIiH, indIiH + 1});
|
||||
REQUIRE_TRUE(gradO->isSameShape(expectedGradOShape), 0,
|
||||
"MAXPOOL2D_BP MKLDNN 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());
|
||||
|
||||
if (paddingMode) // SAME
|
||||
ConvolutionUtils::calcPadding2D(pH, pW, oH, oW, iH, iW, kH, kW, sH, sW, dH, dW);
|
||||
|
||||
onednnUtils::poolingBpONEDNN(input, gradO, gradI, 0, kH, kW, 0, sH, sW, 0, pH, pW, isNCHW, algorithm::pooling_max);
|
||||
|
||||
return sd::Status::OK;
|
||||
}
|
||||
|
||||
PLATFORM_CHECK(maxpool2d_bp, ENGINE_CPU) {
|
||||
auto input = INPUT_VARIABLE(0);
|
||||
auto gradO = INPUT_VARIABLE(1);
|
||||
auto output = OUTPUT_VARIABLE(0);
|
||||
|
||||
Requirements req("ONEDNN MAXPOOL2d_BP OP");
|
||||
req.expectTrue(block.isUseONEDNN(), IS_USE_ONEDNN_MSG) &&
|
||||
req.expectTrue(sd::ONEDNNStream::isSupported({input, output}), ONEDNN_STREAM_NOT_SUPPORTED);
|
||||
if (req) onednnUtils::checkPoolingONEDNN(req, block, input, gradO);
|
||||
req.logTheSuccess();
|
||||
return req;
|
||||
}
|
||||
|
||||
} // namespace platforms
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,156 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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 <helpers/MKLDNNStream.h>
|
||||
#include <ops/declarable/OpRegistrator.h>
|
||||
#include <ops/declarable/PlatformHelper.h>
|
||||
#include <ops/declarable/helpers/convolutions.h>
|
||||
#include <system/platform_boilerplate.h>
|
||||
|
||||
#include "mkldnnUtils.h"
|
||||
|
||||
using namespace dnnl;
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace platforms {
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
PLATFORM_IMPL(maxpool3dnew, ENGINE_CPU) {
|
||||
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)
|
||||
|
||||
sd::LongType kD = INT_ARG(0); // filter(kernel) depth
|
||||
sd::LongType kH = INT_ARG(1); // filter(kernel) height
|
||||
sd::LongType kW = INT_ARG(2); // filter(kernel) width
|
||||
sd::LongType sD = INT_ARG(3); // strides depth
|
||||
sd::LongType sH = INT_ARG(4); // strides height
|
||||
sd::LongType sW = INT_ARG(5); // strides width
|
||||
sd::LongType pD = INT_ARG(6); // paddings depth
|
||||
sd::LongType pH = INT_ARG(7); // paddings height
|
||||
sd::LongType pW = INT_ARG(8); // paddings width
|
||||
sd::LongType dD = INT_ARG(9); // dilations depth
|
||||
sd::LongType dH = INT_ARG(10); // dilations height
|
||||
sd::LongType dW = INT_ARG(11); // dilations width
|
||||
int paddingMode = INT_ARG(12); // 1-SAME, 0-VALID
|
||||
// int extraParam0 = INT_ARG(13); // unnecessary for max case, required only
|
||||
// for avg and pnorm cases
|
||||
int isNCDHW = block.getIArguments()->size() > 14 ? !INT_ARG(14) : 1; // 1-NDHWC, 0-NCDHW
|
||||
|
||||
REQUIRE_TRUE(input->rankOf() == 5, 0,
|
||||
"MAXPOOL3DNEW MKLDNN 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 MKLDNN op: dilation must not be zero, but got instead {%i, %i, %i}", dD, dH, dW);
|
||||
|
||||
sd::LongType bS, iC, iD, iH, iW, oC, oD, oH,
|
||||
oW; // batch size, input channels, input depth/height/width, output channels, output depth/height/width;
|
||||
sd::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);
|
||||
|
||||
if (paddingMode) // SAME
|
||||
ConvolutionUtils::calcPadding3D(pD, pH, pW, oD, oH, oW, iD, iH, iW, kD, kH, kW, sD, sH, sW, dD, dH, dW);
|
||||
|
||||
onednnUtils::poolingONEDNN(input, output, kD, kH, kW, sD, sH, sW, pD, pH, pW, isNCDHW, algorithm::pooling_max);
|
||||
|
||||
return sd::Status::OK;
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
PLATFORM_CHECK(maxpool3dnew, ENGINE_CPU) {
|
||||
auto input = INPUT_VARIABLE(0);
|
||||
auto output = OUTPUT_VARIABLE(0);
|
||||
Requirements req("ONEDNN MAXPOOL3d OP");
|
||||
req.expectTrue(block.isUseONEDNN(), IS_USE_ONEDNN_MSG) &&
|
||||
req.expectTrue(sd::ONEDNNStream::isSupported({input, output}), ONEDNN_STREAM_NOT_SUPPORTED);
|
||||
if (req) onednnUtils::checkPoolingONEDNN(req, block, input, output);
|
||||
req.logTheSuccess();
|
||||
return req;
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
PLATFORM_IMPL(maxpool3dnew_bp, ENGINE_CPU) {
|
||||
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 sd::LongType kD = INT_ARG(0); // filter(kernel) depth
|
||||
const sd::LongType kH = INT_ARG(1); // filter(kernel) height
|
||||
const sd::LongType kW = INT_ARG(2); // filter(kernel) width
|
||||
const sd::LongType sD = INT_ARG(3); // strides depth
|
||||
const sd::LongType sH = INT_ARG(4); // strides height
|
||||
const sd::LongType sW = INT_ARG(5); // strides width
|
||||
sd::LongType pD = INT_ARG(6); // paddings depth
|
||||
sd::LongType pH = INT_ARG(7); // paddings height
|
||||
sd::LongType pW = INT_ARG(8); // paddings width
|
||||
const sd::LongType dD = INT_ARG(9); // dilations depth
|
||||
const sd::LongType dH = INT_ARG(10); // dilations height
|
||||
const sd::LongType dW = INT_ARG(11); // dilations width
|
||||
const int paddngMode = INT_ARG(12); // 1-SAME, 0-VALID
|
||||
// int extraParam0 = INT_ARG(13); // unnecessary for max case, required only
|
||||
// for avg and pnorm cases
|
||||
int isNCDHW = block.getIArguments()->size() > 14 ? !INT_ARG(14) : 1; // 1-NDHWC, 0-NCDHW
|
||||
|
||||
REQUIRE_TRUE(input->rankOf() == 5, 0, "MAXPOOL3DNEW_BP MKLDNN op: input should have rank of 5, but got %i instead",
|
||||
input->rankOf());
|
||||
REQUIRE_TRUE(dD != 0 && dH != 0 && dW != 0, 0,
|
||||
"MAXPOOL3DNEW_BP MKLDNN op: dilation must not be zero, but got instead {%i, %i, %i}", dD, dH, dW);
|
||||
|
||||
sd::LongType bS, iC, iD, iH, iW, oC, oD, oH,
|
||||
oW; // batch size, input channels, input depth/height/width, output channels, output depth/height/width;
|
||||
sd::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<sd::LongType> expectedGradOShape =
|
||||
ShapeUtils::composeShapeUsingDimsAndIdx({bS, iC, oD, oH, oW, 0, indIOioC, indIOioD, indIOioD + 1, indIOioD + 2});
|
||||
REQUIRE_TRUE(gradO->isSameShape(expectedGradOShape), 0,
|
||||
"MAXPOOL3DNEW_BP MKLDNN 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());
|
||||
|
||||
if (paddngMode) // SAME
|
||||
ConvolutionUtils::calcPadding3D(pD, pH, pW, oD, oH, oW, iD, iH, iW, kD, kH, kW, sD, sH, sW, dD, dH, dW);
|
||||
|
||||
onednnUtils::poolingBpONEDNN(input, gradO, gradI, kD, kH, kW, sD, sH, sW, pD, pH, pW, isNCDHW,
|
||||
algorithm::pooling_max);
|
||||
|
||||
return sd::Status::OK;
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
PLATFORM_CHECK(maxpool3dnew_bp, ENGINE_CPU) {
|
||||
auto input = INPUT_VARIABLE(0);
|
||||
auto gradO = INPUT_VARIABLE(1);
|
||||
auto output = OUTPUT_VARIABLE(0);
|
||||
|
||||
Requirements req("ONEDNN MAXPOOL3d_BP OP");
|
||||
req.expectTrue(block.isUseONEDNN(), IS_USE_ONEDNN_MSG) &&
|
||||
req.expectTrue(sd::ONEDNNStream::isSupported({input, output}), ONEDNN_STREAM_NOT_SUPPORTED);
|
||||
if (req) onednnUtils::checkPoolingONEDNN(req, block, input, gradO);
|
||||
req.logTheSuccess();
|
||||
return req;
|
||||
}
|
||||
|
||||
} // namespace platforms
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,440 @@
|
||||
/* ******************************************************************************
|
||||
*
|
||||
*
|
||||
* 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 saudet
|
||||
// @author Yurii Shyrma (iuriish@yahoo.com)
|
||||
//
|
||||
#include "mkldnnUtils.h"
|
||||
|
||||
#include <dnnl_types.h>
|
||||
#include <ops/declarable/helpers/convolutions.h>
|
||||
|
||||
using namespace dnnl;
|
||||
|
||||
namespace sd {
|
||||
namespace onednnUtils {
|
||||
|
||||
//////////////////////////////////////////////////////////////////////
|
||||
void getDims(NDArray* array, const int rank, dnnl::memory::dims& mklDims) {
|
||||
std::vector<int64_t> vDims(rank);
|
||||
for (auto i = 0; i < rank; i++) {
|
||||
vDims[i] = array->sizeAt(i);
|
||||
}
|
||||
mklDims = dnnl::memory::dims(vDims);
|
||||
}
|
||||
//////////////////////////////////////////////////////////////////////
|
||||
dnnl::memory::format_tag getFormat(NDArray& arr) {
|
||||
dnnl::memory::format_tag result;
|
||||
|
||||
switch (arr.rankOf()) {
|
||||
case 1:
|
||||
result = dnnl::memory::format_tag::a;
|
||||
break;
|
||||
case 2:
|
||||
result = arr.ordering() == 'c' ? dnnl::memory::format_tag::ab : dnnl::memory::format_tag::ba;
|
||||
break;
|
||||
case 3:
|
||||
result = arr.ordering() == 'c' ? dnnl::memory::format_tag::abc : dnnl::memory::format_tag::cba;
|
||||
break;
|
||||
case 4:
|
||||
result = dnnl::memory::format_tag::abcd;
|
||||
break;
|
||||
case 5:
|
||||
result = dnnl::memory::format_tag::abcde;
|
||||
break;
|
||||
case 6:
|
||||
result = dnnl::memory::format_tag::abcdef;
|
||||
break;
|
||||
default:
|
||||
THROW_EXCEPTION("MKLDNN getFormat: do we really want to use arras with rank > 6 ?");
|
||||
}
|
||||
|
||||
return result;
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////
|
||||
void setBlockStrides(NDArray& array, dnnl::memory::desc& mklMd, const std::vector<int>& permut) {
|
||||
if ((array.rankOf() > 3 && array.ordering() == 'f') || !permut.empty()) {
|
||||
mklMd.data.format_kind = dnnl_blocked; // overrides format
|
||||
|
||||
if (permut.empty())
|
||||
for (auto i = 0; i < array.rankOf(); ++i) mklMd.data.format_desc.blocking.strides[i] = array.strideAt(i);
|
||||
else {
|
||||
if (static_cast<size_t>(array.rankOf()) != permut.size())
|
||||
THROW_EXCEPTION("mkldnnUtils::setBlockStrides: size of permut vector is not equal to array rank !");
|
||||
for (auto i = 0; i < array.rankOf(); ++i) mklMd.data.format_desc.blocking.strides[i] = array.strideAt(permut[i]);
|
||||
}
|
||||
}
|
||||
}
|
||||
////////////////////////////////////////////////////////////////////////////////////////////////
|
||||
dnnl::memory loadDataToMklStream(NDArray& array, const dnnl::engine& engine, const dnnl::stream& stream,
|
||||
const dnnl::memory::desc& user_md, const dnnl::memory::desc& primitive_md,
|
||||
dnnl::memory& arg) {
|
||||
auto user_mem = dnnl::memory(user_md, engine, const_cast<NDArray&>(array).buffer());
|
||||
const bool bReorder = primitive_md != user_mem.get_desc();
|
||||
auto mkl_mem = bReorder ? dnnl::memory(primitive_md, engine) : user_mem;
|
||||
if (bReorder) dnnl::reorder(user_mem, mkl_mem).execute(stream, user_mem, mkl_mem);
|
||||
arg = mkl_mem;
|
||||
return user_mem;
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////
|
||||
void poolingONEDNN(NDArray* input, NDArray* output, const sd::LongType kD, const sd::LongType kH, const sd::LongType kW, const sd::LongType sD,
|
||||
const sd::LongType sH, const sd::LongType sW, const sd::LongType pD, const sd::LongType pH, const sd::LongType pW, const int isNCHW,
|
||||
const dnnl::algorithm mode) {
|
||||
// unfortunately mkl dnn doesn't support any format (dnnl::memory::format_tag::any) for input
|
||||
const sd::LongType rank = input->rankOf();
|
||||
|
||||
sd::LongType bS, iC, iD, iH, iW, oC, oD, oH, oW, indIOioC, indIiH, indWoC, indWiC, indWkH, indOoH;
|
||||
dnnl::memory::dims strides, kernel, padding, padding_r, xDims, zDims;
|
||||
dnnl::memory::format_tag xzFrmat;
|
||||
|
||||
const auto type = dnnl::memory::data_type::f32;
|
||||
|
||||
if (rank == 4) { // 2d
|
||||
|
||||
ops::ConvolutionUtils::getSizesAndIndexesConv2d(isNCHW, 0, *input, *output, bS, iC, iH, iW, oC, oH, oW, indIOioC,
|
||||
indIiH, indWiC, indWoC, indWkH, indOoH);
|
||||
|
||||
strides = {sH, sW};
|
||||
kernel = {kH, kW};
|
||||
padding = {pH, pW};
|
||||
padding_r = {(oH - 1) * sH - iH + kH - pH, (oW - 1) * sW - iW + kW - pW};
|
||||
xDims = {bS, iC, iH, iW};
|
||||
zDims = {bS, oC, oH, oW};
|
||||
|
||||
xzFrmat = isNCHW ? dnnl::memory::format_tag::nchw : dnnl::memory::format_tag::nhwc;
|
||||
} else { // 3d
|
||||
|
||||
ops::ConvolutionUtils::getSizesAndIndexesConv3d(isNCHW, 0, *input, *output, bS, iC, iD, iH, iW, oC, oD, oH, oW,
|
||||
indIOioC, indIiH, indWiC, indWoC, indWkH);
|
||||
|
||||
strides = {sD, sH, sW};
|
||||
kernel = {kD, kH, kW};
|
||||
padding = {pD, pH, pW};
|
||||
padding_r = {(oD - 1) * sD - iD + kD - pD, (oH - 1) * sH - iH + kH - pH, (oW - 1) * sW - iW + kW - pW};
|
||||
xDims = {bS, iC, iD, iH, iW};
|
||||
zDims = {bS, oC, oD, oH, oW};
|
||||
|
||||
xzFrmat = isNCHW ? dnnl::memory::format_tag::ncdhw : dnnl::memory::format_tag::ndhwc;
|
||||
}
|
||||
|
||||
std::vector<int> permut;
|
||||
if (!isNCHW) permut = rank == 4 ? std::vector<int>({0, 3, 1, 2}) : std::vector<int>({0, 4, 1, 2, 3});
|
||||
|
||||
// memory descriptors for arrays
|
||||
|
||||
// input
|
||||
dnnl::memory::desc x_mkl_md = dnnl::memory::desc(xDims, type, xzFrmat);
|
||||
dnnl::memory::desc x_user_md = dnnl::memory::desc(xDims, type, xzFrmat);
|
||||
onednnUtils::setBlockStrides(*input, x_user_md, permut);
|
||||
|
||||
// output
|
||||
dnnl::memory::desc z_mkl_md = dnnl::memory::desc(zDims, type, dnnl::memory::format_tag::any);
|
||||
dnnl::memory::desc z_user_md = dnnl::memory::desc(zDims, type, xzFrmat);
|
||||
onednnUtils::setBlockStrides(*output, z_user_md, permut);
|
||||
|
||||
auto engine = onednnUtils::getEngine(LaunchContext::defaultContext()->engine());
|
||||
|
||||
// operation primitive description
|
||||
dnnl::pooling_forward::desc op_desc(dnnl::prop_kind::forward_inference, mode, x_mkl_md, z_mkl_md, strides, kernel,
|
||||
padding, padding_r);
|
||||
dnnl::pooling_forward::primitive_desc op_prim_desc(op_desc, engine);
|
||||
|
||||
// arguments (memory buffers) necessary for calculations
|
||||
std::unordered_map<int, dnnl::memory> args;
|
||||
|
||||
dnnl::stream stream(engine);
|
||||
|
||||
// provide memory buffers and check whether reorder is required
|
||||
|
||||
// input
|
||||
onednnUtils::loadDataToMklStream(*input, engine, stream, x_user_md, op_prim_desc.src_desc(), args[DNNL_ARG_SRC]);
|
||||
|
||||
// output
|
||||
auto z_user_mem =
|
||||
onednnUtils::loadDataToMklStream(*output, engine, stream, z_user_md, op_prim_desc.dst_desc(), args[DNNL_ARG_DST]);
|
||||
|
||||
// run calculations
|
||||
dnnl::pooling_forward(op_prim_desc).execute(stream, args);
|
||||
|
||||
// reorder outputs if necessary
|
||||
if (op_prim_desc.dst_desc() != z_user_mem.get_desc())
|
||||
dnnl::reorder(args[DNNL_ARG_DST], z_user_mem).execute(stream, args[DNNL_ARG_DST], z_user_mem);
|
||||
|
||||
stream.wait();
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////
|
||||
void poolingBpONEDNN(NDArray* input, NDArray* gradO, NDArray* gradI, const sd::LongType kD, const sd::LongType kH,
|
||||
const sd::LongType kW, const sd::LongType sD, const sd::LongType sH, const sd::LongType sW, const sd::LongType pD, const sd::LongType pH, const sd::LongType pW,
|
||||
const int isNCHW, const dnnl::algorithm mode) {
|
||||
// unfortunately mkl dnn doesn't support any format (dnnl::memory::format_tag::any) for input
|
||||
|
||||
const sd::LongType rank = input->rankOf();
|
||||
|
||||
sd::LongType bS, iC, iD, iH, iW, oC, oD, oH, oW, indIOioC, indIiH, indWoC, indWiC, indWkH, indOoH;
|
||||
dnnl::memory::dims strides, kernel, padding, padding_r, xDims, zDims;
|
||||
dnnl::memory::format_tag xzFrmat;
|
||||
|
||||
const auto type = dnnl::memory::data_type::f32;
|
||||
|
||||
if (rank == 4) { // 2d
|
||||
|
||||
ops::ConvolutionUtils::getSizesAndIndexesConv2d(isNCHW, 0, *input, *gradO, bS, iC, iH, iW, oC, oH, oW, indIOioC,
|
||||
indIiH, indWiC, indWoC, indWkH, indOoH);
|
||||
|
||||
strides = {sH, sW};
|
||||
kernel = {kH, kW};
|
||||
padding = {pH, pW};
|
||||
padding_r = {(oH - 1) * sH - iH + kH - pH, (oW - 1) * sW - iW + kW - pW};
|
||||
xDims = {bS, iC, iH, iW};
|
||||
zDims = {bS, oC, oH, oW};
|
||||
|
||||
xzFrmat = isNCHW ? dnnl::memory::format_tag::nchw : dnnl::memory::format_tag::nhwc;
|
||||
} else { // 3d
|
||||
|
||||
ops::ConvolutionUtils::getSizesAndIndexesConv3d(isNCHW, 0, *input, *gradO, bS, iC, iD, iH, iW, oC, oD, oH, oW,
|
||||
indIOioC, indIiH, indWiC, indWoC, indWkH);
|
||||
|
||||
strides = {sD, sH, sW};
|
||||
kernel = {kD, kH, kW};
|
||||
padding = {pD, pH, pW};
|
||||
padding_r = {(oD - 1) * sD - iD + kD - pD, (oH - 1) * sH - iH + kH - pH, (oW - 1) * sW - iW + kW - pW};
|
||||
xDims = {bS, iC, iD, iH, iW};
|
||||
zDims = {bS, oC, oD, oH, oW};
|
||||
|
||||
xzFrmat = isNCHW ? dnnl::memory::format_tag::ncdhw : dnnl::memory::format_tag::ndhwc;
|
||||
}
|
||||
|
||||
std::vector<int> permut;
|
||||
if (!isNCHW) permut = rank == 4 ? std::vector<int>({0, 3, 1, 2}) : std::vector<int>({0, 4, 1, 2, 3});
|
||||
|
||||
// memory descriptors for arrays
|
||||
|
||||
// input
|
||||
dnnl::memory::desc x_mkl_md = dnnl::memory::desc(xDims, type, xzFrmat);
|
||||
dnnl::memory::desc x_user_md = dnnl::memory::desc(xDims, type, xzFrmat);
|
||||
onednnUtils::setBlockStrides(*input, x_user_md, permut);
|
||||
|
||||
// gradO
|
||||
dnnl::memory::desc gradO_mkl_md = dnnl::memory::desc(zDims, type, dnnl::memory::format_tag::any);
|
||||
dnnl::memory::desc gradO_user_md = dnnl::memory::desc(zDims, type, xzFrmat);
|
||||
onednnUtils::setBlockStrides(*gradO, gradO_user_md, permut);
|
||||
|
||||
// gradI
|
||||
dnnl::memory::desc gradI_mkl_md = dnnl::memory::desc(xDims, type, dnnl::memory::format_tag::any);
|
||||
dnnl::memory::desc gradI_user_md = dnnl::memory::desc(xDims, type, xzFrmat);
|
||||
onednnUtils::setBlockStrides(*gradI, gradI_user_md, permut);
|
||||
|
||||
auto engine = onednnUtils::getEngine(LaunchContext::defaultContext()->engine());
|
||||
dnnl::stream stream(engine);
|
||||
|
||||
// forward primitive description
|
||||
dnnl::pooling_forward::desc op_ff_desc(dnnl::prop_kind::forward, mode, x_mkl_md, gradO_mkl_md, strides, kernel,
|
||||
padding, padding_r);
|
||||
dnnl::pooling_forward::primitive_desc op_ff_prim_desc(op_ff_desc, engine);
|
||||
|
||||
// backward primitive description
|
||||
dnnl::pooling_backward::desc op_bp_desc(mode, gradI_mkl_md, gradO_mkl_md, strides, kernel, padding, padding_r);
|
||||
dnnl::pooling_backward::primitive_desc op_bp_prim_desc(op_bp_desc, engine, op_ff_prim_desc);
|
||||
|
||||
// arguments (memory buffers) necessary for calculations
|
||||
std::unordered_map<int, dnnl::memory> args;
|
||||
|
||||
// gradO
|
||||
onednnUtils::loadDataToMklStream(*gradO, engine, stream, gradO_user_md, op_bp_prim_desc.diff_dst_desc(),
|
||||
args[DNNL_ARG_DIFF_DST]);
|
||||
|
||||
// gradI
|
||||
auto gradI_user_mem = onednnUtils::loadDataToMklStream(*gradI, engine, stream, gradI_user_md,
|
||||
op_bp_prim_desc.diff_src_desc(), args[DNNL_ARG_DIFF_SRC]);
|
||||
|
||||
if (mode == algorithm::pooling_max) {
|
||||
// input
|
||||
onednnUtils::loadDataToMklStream(*input, engine, stream, x_user_md, op_ff_prim_desc.src_desc(), args[DNNL_ARG_SRC]);
|
||||
|
||||
// z
|
||||
auto z_mkl_mem = dnnl::memory(op_ff_prim_desc.dst_desc(), engine);
|
||||
args[DNNL_ARG_DST] = z_mkl_mem;
|
||||
|
||||
// auxiliary memory allocation
|
||||
auto workspace = dnnl::memory(op_ff_prim_desc.workspace_desc(), engine);
|
||||
args[DNNL_ARG_WORKSPACE] = workspace;
|
||||
|
||||
// run forward calculations
|
||||
dnnl::pooling_forward(op_ff_prim_desc).execute(stream, args);
|
||||
}
|
||||
|
||||
// run backward calculations
|
||||
dnnl::pooling_backward(op_bp_prim_desc).execute(stream, args);
|
||||
|
||||
// reorder gradI if necessary
|
||||
if (op_bp_prim_desc.diff_src_desc() != gradI_user_mem.get_desc())
|
||||
dnnl::reorder(args[DNNL_ARG_DIFF_SRC], gradI_user_mem).execute(stream, args[DNNL_ARG_DIFF_SRC], gradI_user_mem);
|
||||
|
||||
stream.wait();
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
void getONEDNNMemoryDescLrn(NDArray* src, NDArray* diff_src, NDArray* dst,
|
||||
dnnl::memory::desc* lrn_src_md, dnnl::memory::desc* lrn_diff_src_md,
|
||||
dnnl::memory::desc* lrn_dst_md, dnnl::memory::desc* user_src_md,
|
||||
dnnl::memory::desc* user_diff_src_md, dnnl::memory::desc* user_dst_md, int axis) {
|
||||
const sd::LongType* shape = src->shapeInfo();
|
||||
long rank = shape[0];
|
||||
long dim1 = axis; // MKL-DNN supports only 1 axis, which has to be the "channel" one
|
||||
long dim2 = axis >= 2 ? 1 : 2;
|
||||
long dim3 = axis >= 3 ? 2 : 3;
|
||||
dnnl::memory::dims lrn_src_tz = {(int)shape[1], (int)shape[dim1 + 1], rank > 2 ? (int)shape[dim2 + 1] : 1,
|
||||
rank > 3 ? (int)shape[dim3 + 1] : 1};
|
||||
|
||||
auto type = dnnl::memory::data_type::f32;
|
||||
auto format = axis == 1 ? dnnl::memory::format_tag::nchw : dnnl::memory::format_tag::nhwc;
|
||||
auto supposed_to_be_any_format = format; // doesn't work with "any"
|
||||
|
||||
if (src != nullptr && src->buffer() != nullptr && lrn_src_md != nullptr) {
|
||||
*lrn_src_md = dnnl::memory::desc({lrn_src_tz}, type, supposed_to_be_any_format);
|
||||
*user_src_md = dnnl::memory::desc({lrn_src_tz}, type, format);
|
||||
user_src_md->data.format_kind = dnnl_blocked;
|
||||
user_src_md->data.format_desc.blocking.strides[0] = src->stridesOf()[0];
|
||||
user_src_md->data.format_desc.blocking.strides[1] = src->stridesOf()[dim1];
|
||||
user_src_md->data.format_desc.blocking.strides[2] = rank > 2 ? src->stridesOf()[dim2] : 1;
|
||||
user_src_md->data.format_desc.blocking.strides[3] = rank > 3 ? src->stridesOf()[dim3] : 1;
|
||||
}
|
||||
|
||||
if (diff_src != nullptr && diff_src->buffer() != nullptr && lrn_diff_src_md != nullptr) {
|
||||
*lrn_diff_src_md = dnnl::memory::desc({lrn_src_tz}, type, supposed_to_be_any_format);
|
||||
*user_diff_src_md = dnnl::memory::desc({lrn_src_tz}, type, format);
|
||||
user_diff_src_md->data.format_kind = dnnl_blocked;
|
||||
user_diff_src_md->data.format_desc.blocking.strides[0] = diff_src->stridesOf()[0];
|
||||
user_diff_src_md->data.format_desc.blocking.strides[1] = diff_src->stridesOf()[dim1];
|
||||
user_diff_src_md->data.format_desc.blocking.strides[2] = rank > 2 ? diff_src->stridesOf()[dim2] : 1;
|
||||
user_diff_src_md->data.format_desc.blocking.strides[3] = rank > 3 ? diff_src->stridesOf()[dim3] : 1;
|
||||
}
|
||||
|
||||
if (dst != nullptr && dst->buffer() != nullptr && lrn_dst_md != nullptr) {
|
||||
*lrn_dst_md = dnnl::memory::desc({lrn_src_tz}, type, supposed_to_be_any_format);
|
||||
*user_dst_md = dnnl::memory::desc({lrn_src_tz}, type, format);
|
||||
user_dst_md->data.format_kind = dnnl_blocked;
|
||||
user_dst_md->data.format_desc.blocking.strides[0] = dst->stridesOf()[0];
|
||||
user_dst_md->data.format_desc.blocking.strides[1] = dst->stridesOf()[dim1];
|
||||
user_dst_md->data.format_desc.blocking.strides[2] = rank > 2 ? dst->stridesOf()[dim2] : 1;
|
||||
user_dst_md->data.format_desc.blocking.strides[3] = rank > 3 ? dst->stridesOf()[dim3] : 1;
|
||||
}
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////////
|
||||
dnnl::engine& getEngine(void* ptr) {
|
||||
auto eng = reinterpret_cast<dnnl::engine*>(ptr);
|
||||
return *eng;
|
||||
}
|
||||
|
||||
void checkPoolingONEDNN(Requirements& reqs, sd::graph::Context& block, sd::NDArray* in, sd::NDArray* out) {
|
||||
// replicate OneDNN check that was added since v1.8
|
||||
// https://github.com/oneapi-src/oneDNN/blob/master/src/common/pooling.cpp#L108-L110
|
||||
// if (str < 1 || dil < 0 || pad_l < 0 || pad_r + str < 0) return invalid_arguments;
|
||||
if (in->rankOf() > 4 && block.getIArguments()->size() > 12) {
|
||||
// pooling 3D
|
||||
sd::LongType kD = INT_ARG(0); // filter(kernel) depth
|
||||
sd::LongType kH = INT_ARG(1); // filter(kernel) height
|
||||
sd::LongType kW = INT_ARG(2); // filter(kernel) width
|
||||
sd::LongType sD = INT_ARG(3); // strides depth
|
||||
sd::LongType sH = INT_ARG(4); // strides height
|
||||
sd::LongType sW = INT_ARG(5); // strides width
|
||||
sd::LongType pD = INT_ARG(6); // paddings depth
|
||||
sd::LongType pH = INT_ARG(7); // paddings height
|
||||
sd::LongType pW = INT_ARG(8); // paddings width
|
||||
sd::LongType dD = INT_ARG(9); // dilations depth
|
||||
sd::LongType dH = INT_ARG(10); // dilations height
|
||||
sd::LongType dW = INT_ARG(11); // dilations width
|
||||
sd::LongType paddingMode = INT_ARG(12); // 1-SAME, 0-VALID
|
||||
// int extraParam0 = INT_ARG(13); // unnecessary for max case, required only for avg and pnorm cases
|
||||
int isNCDHW = block.getIArguments()->size() > 14 ? !INT_ARG(14) : 1; // 1-NDHWC, 0-NCDHW
|
||||
reqs.expectEq(makeInfoVariable(in->rankOf(), RANK_MSG_INPUT0), 5) &&
|
||||
// stride >=1
|
||||
reqs.expectGreaterEq(makeInfoVariable(sD, "strides#Depth"), 1) &&
|
||||
reqs.expectGreaterEq(makeInfoVariable(sH, "strides#Height"), 1) &&
|
||||
reqs.expectGreaterEq(makeInfoVariable(sW, "strides#Width"), 1) &&
|
||||
// dilation >=0
|
||||
reqs.expectGreaterEq(makeInfoVariable(dW, "dilation#Depth"), 0) &&
|
||||
reqs.expectGreaterEq(makeInfoVariable(dH, "dilation#Height"), 0) &&
|
||||
reqs.expectGreaterEq(makeInfoVariable(dW, "dilation#Width"), 0);
|
||||
if (reqs) {
|
||||
sd::LongType bS, iC, iD, iH, iW, oC, oD, oH,
|
||||
oW; // batch size, input channels, input depth/height/width, output channels, output depth/height/width;
|
||||
sd::LongType indIOioC, indIOioD, indWoC, indWiC, indWkD; // corresponding indexes
|
||||
ops::ConvolutionUtils::getSizesAndIndexesConv3d(isNCDHW, 0, *in, *out, bS, iC, iD, iH, iW, oC, oD, oH, oW, indIOioC,
|
||||
indIOioD, indWiC, indWoC, indWkD);
|
||||
|
||||
if (paddingMode) // SAME
|
||||
ops::ConvolutionUtils::calcPadding3D(pD, pH, pW, oD, oH, oW, iD, iH, iW, kD, kH, kW, sD, sH, sW, dD, dH, dW);
|
||||
// pad_l >=0
|
||||
reqs.expectGreaterEq(makeInfoVariable(pD, "padding_l#Depth"), 0) &&
|
||||
reqs.expectGreaterEq(makeInfoVariable(pH, "padding_l#Height"), 0) &&
|
||||
reqs.expectGreaterEq(makeInfoVariable(pW, "padding_l#Width"), 0) &&
|
||||
// pad_r+ stride
|
||||
reqs.expectGreaterEq(makeInfoVariable(((oD - 1) * sD - iD + kD - pD) + sD, "padding_r#Depth + stride#Depth"),
|
||||
0) &&
|
||||
reqs.expectGreaterEq(
|
||||
makeInfoVariable(((oH - 1) * sH - iH + kH - pH) + sH, "padding_r#Height + stride#Height"), 0) &&
|
||||
reqs.expectGreaterEq(makeInfoVariable(((oW - 1) * sW - iW + kW - pW) + sW, "padding_r#Width + stride#Width"),
|
||||
0);
|
||||
}
|
||||
} else if (block.getIArguments()->size() > 8) {
|
||||
const int kH = INT_ARG(0);
|
||||
const int kW = INT_ARG(1);
|
||||
const int sH = INT_ARG(2);
|
||||
const int sW = INT_ARG(3);
|
||||
sd::LongType pH = INT_ARG(4);
|
||||
sd::LongType pW = INT_ARG(5);
|
||||
const int dH = INT_ARG(6);
|
||||
const int dW = INT_ARG(7);
|
||||
const int paddingMode = INT_ARG(8);
|
||||
const int isNCHW = block.getIArguments()->size() > 10 ? !INT_ARG(10) : 1; // INT_ARG(10): 1-NHWC, 0-NCHW
|
||||
reqs.expectEq(makeInfoVariable(in->rankOf(), RANK_MSG_INPUT0), 4) &&
|
||||
// stride >=1
|
||||
reqs.expectGreaterEq(makeInfoVariable(sH, "strides#Height"), 1) &&
|
||||
reqs.expectGreaterEq(makeInfoVariable(sW, "strides#Width"), 1) &&
|
||||
// dilation >=0
|
||||
reqs.expectGreaterEq(makeInfoVariable(dH, "dilation#Height"), 0) &&
|
||||
reqs.expectGreaterEq(makeInfoVariable(dW, "dilation#Width"), 0);
|
||||
if (reqs) {
|
||||
sd::LongType bS, iC, iD, iH, iW, oC, oD, oH, oW, indIOioC, indIiH, indWoC, indWiC, indWkH, indOoH;
|
||||
ops::ConvolutionUtils::getSizesAndIndexesConv2d(isNCHW, 0, *in, *out, bS, iC, iH, iW, oC, oH, oW, indIOioC, indIiH,
|
||||
indWiC, indWoC, indWkH, indOoH);
|
||||
if (paddingMode) {
|
||||
ops::ConvolutionUtils::calcPadding2D(pH, pW, oH, oW, iH, iW, kH, kW, sH, sW, dH, dW);
|
||||
}
|
||||
// pad_l >=0
|
||||
reqs.expectGreaterEq(makeInfoVariable(pH, "padding_l#Height"), 0) &&
|
||||
reqs.expectGreaterEq(makeInfoVariable(pW, "padding_l#Width"), 0) &&
|
||||
// pad_r+ stride
|
||||
reqs.expectGreaterEq(
|
||||
makeInfoVariable(((oH - 1) * sH - iH + kH - pH) + sH, "padding_r#Height + stride#Height"), 0) &&
|
||||
reqs.expectGreaterEq(makeInfoVariable(((oW - 1) * sW - iW + kW - pW) + sW, "padding_r#Width + stride#Width"),
|
||||
0);
|
||||
}
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
|
||||
} // namespace onednnUtils
|
||||
} // 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 saudet
|
||||
// @author Yurii Shyrma (iuriish@yahoo.com)
|
||||
//
|
||||
|
||||
#ifndef DEV_TESTS_MKLDNNUTILS_H
|
||||
#define DEV_TESTS_MKLDNNUTILS_H
|
||||
|
||||
#include <array/NDArray.h>
|
||||
#include <graph/Context.h>
|
||||
#include <helpers/MKLDNNStream.h>
|
||||
#include <legacy/NativeOps.h>
|
||||
#include <ops/declarable/PlatformHelper.h>
|
||||
#include <system/platform_boilerplate.h>
|
||||
|
||||
#include <dnnl.hpp>
|
||||
|
||||
using namespace samediff;
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace platforms {
|
||||
/**
|
||||
* Here we actually declare our platform helpers
|
||||
*/
|
||||
DECLARE_PLATFORM(conv2d, ENGINE_CPU);
|
||||
|
||||
DECLARE_PLATFORM(conv2d_bp, ENGINE_CPU);
|
||||
|
||||
DECLARE_PLATFORM(avgpool2d, ENGINE_CPU);
|
||||
|
||||
DECLARE_PLATFORM(avgpool2d_bp, ENGINE_CPU);
|
||||
|
||||
DECLARE_PLATFORM(maxpool2d, ENGINE_CPU);
|
||||
|
||||
DECLARE_PLATFORM(maxpool2d_bp, ENGINE_CPU);
|
||||
|
||||
DECLARE_PLATFORM(conv3dnew, ENGINE_CPU);
|
||||
|
||||
DECLARE_PLATFORM(conv3dnew_bp, ENGINE_CPU);
|
||||
|
||||
DECLARE_PLATFORM(maxpool3dnew, ENGINE_CPU);
|
||||
|
||||
DECLARE_PLATFORM(maxpool3dnew_bp, ENGINE_CPU);
|
||||
|
||||
DECLARE_PLATFORM(avgpool3dnew, ENGINE_CPU);
|
||||
|
||||
DECLARE_PLATFORM(avgpool3dnew_bp, ENGINE_CPU);
|
||||
|
||||
DECLARE_PLATFORM(lrn, ENGINE_CPU);
|
||||
|
||||
DECLARE_PLATFORM(batchnorm, ENGINE_CPU);
|
||||
|
||||
DECLARE_PLATFORM(batchnorm_bp, ENGINE_CPU);
|
||||
|
||||
DECLARE_PLATFORM(lstmLayer, ENGINE_CPU);
|
||||
|
||||
DECLARE_PLATFORM(deconv2d, ENGINE_CPU);
|
||||
|
||||
DECLARE_PLATFORM(deconv2d_tf, ENGINE_CPU);
|
||||
|
||||
DECLARE_PLATFORM(deconv3d, ENGINE_CPU);
|
||||
|
||||
DECLARE_PLATFORM(deconv2d_bp, ENGINE_CPU);
|
||||
|
||||
DECLARE_PLATFORM(deconv3d_bp, ENGINE_CPU);
|
||||
|
||||
DECLARE_PLATFORM(depthwise_conv2d, ENGINE_CPU);
|
||||
|
||||
DECLARE_PLATFORM(depthwise_conv2d_bp, ENGINE_CPU);
|
||||
|
||||
DECLARE_PLATFORM(matmul, ENGINE_CPU);
|
||||
|
||||
DECLARE_PLATFORM(softmax, ENGINE_CPU);
|
||||
|
||||
DECLARE_PLATFORM(softmax_bp, ENGINE_CPU);
|
||||
|
||||
DECLARE_PLATFORM(tanh, ENGINE_CPU);
|
||||
|
||||
DECLARE_PLATFORM(tanh_bp, ENGINE_CPU);
|
||||
|
||||
DECLARE_PLATFORM(xw_plus_b, ENGINE_CPU);
|
||||
|
||||
DECLARE_PLATFORM(xw_plus_b_bp, ENGINE_CPU);
|
||||
|
||||
DECLARE_PLATFORM(concat, ENGINE_CPU);
|
||||
|
||||
} // namespace platforms
|
||||
} // namespace ops
|
||||
|
||||
namespace onednnUtils {
|
||||
|
||||
void poolingONEDNN(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 isNCHW,
|
||||
const dnnl::algorithm mode);
|
||||
|
||||
void poolingBpONEDNN(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 isNCHW, const dnnl::algorithm mode);
|
||||
|
||||
void getONEDNNMemoryDescLrn(NDArray* src, NDArray* diff_src, NDArray* dst,
|
||||
dnnl::memory::desc* lrn_src_md, dnnl::memory::desc* lrn_diff_src_md,
|
||||
dnnl::memory::desc* lrn_dst_md, dnnl::memory::desc* user_src_md,
|
||||
dnnl::memory::desc* user_diff_src_md, dnnl::memory::desc* user_dst_md, int axis);
|
||||
|
||||
dnnl::engine& getEngine(void* ptr);
|
||||
|
||||
/**
|
||||
* This function creates memory dimentions
|
||||
* @param const pointer to array
|
||||
* @param const array rank
|
||||
* @param reference to memory dimentions
|
||||
*/
|
||||
void getDims(NDArray* array, const int rank, dnnl::memory::dims& mklDims);
|
||||
/**
|
||||
* This function evaluate memory format tag based on array shapeInfo
|
||||
* @param const array
|
||||
* @return memory format
|
||||
*/
|
||||
dnnl::memory::format_tag getFormat(NDArray& arr);
|
||||
|
||||
void setBlockStrides(NDArray& array, dnnl::memory::desc& mklMd, const std::vector<int>& permut = {});
|
||||
//////////////////////////////////////////////////////////////////////
|
||||
/**
|
||||
* This function load and reorder user memory to mkl
|
||||
* @param const pointer to dataset
|
||||
* @param reference to mkl engine
|
||||
* @param reference to mkl stream
|
||||
* @param reference to args container for dnnl
|
||||
* @param reference to user memory description
|
||||
* @param primitive memory descriptor
|
||||
* @param dnnl arg activation enumerator
|
||||
*/
|
||||
dnnl::memory loadDataToMklStream(NDArray& array, const dnnl::engine& engine, const dnnl::stream& stream,
|
||||
const dnnl::memory::desc& user_md, const dnnl::memory::desc& primitive_md,
|
||||
dnnl::memory& arg);
|
||||
|
||||
/**
|
||||
* @brief This function checks adittional ONEDNN pooling requirements
|
||||
*
|
||||
* @param reqs Requirements block to store the check result
|
||||
* @param block Context block to extract positional integer arguments.
|
||||
* @param in in NDArray
|
||||
* @param out out NDArray
|
||||
*/
|
||||
void checkPoolingONEDNN(Requirements& reqs, sd::graph::Context& block, sd::NDArray* in, sd::NDArray* out);
|
||||
|
||||
|
||||
|
||||
} // namespace onednnUtils
|
||||
} // namespace sd
|
||||
|
||||
#endif // DEV_TESTS_MKLDNNUTILS_H
|
||||
@@ -0,0 +1,268 @@
|
||||
/*
|
||||
* ******************************************************************************
|
||||
* *
|
||||
* *
|
||||
* * 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 Oleg Semeniv <oleg.semeniv@gmail.com>
|
||||
//
|
||||
//
|
||||
#include <helpers/MKLDNNStream.h>
|
||||
#include <ops/declarable/OpRegistrator.h>
|
||||
#include <ops/declarable/PlatformHelper.h>
|
||||
#include <system/platform_boilerplate.h>
|
||||
|
||||
#include "mkldnnUtils.h"
|
||||
|
||||
using namespace dnnl;
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace platforms {
|
||||
|
||||
//////////////////////////////////////////////////////////////////////
|
||||
static void softmaxMKLDNN(NDArray* x, NDArray* z, const sd::LongType axis) {
|
||||
dnnl::memory::dims shape = x->getShapeAsFlatVector();
|
||||
|
||||
const sd::LongType xRank = x->rankOf();
|
||||
|
||||
dnnl::memory::format_tag xFormat = onednnUtils::getFormat(*x);
|
||||
dnnl::memory::format_tag zFormat = onednnUtils::getFormat(*z);
|
||||
|
||||
// optimized cases
|
||||
if (2 == xRank && 0 == axis) {
|
||||
xFormat = dnnl::memory::format_tag::ba;
|
||||
zFormat = dnnl::memory::format_tag::ba;
|
||||
} else if (4 == xRank && 1 == axis && (x->sizeAt(2) * x->sizeAt(3)) > 1) {
|
||||
xFormat = dnnl::memory::format_tag::acdb;
|
||||
zFormat = dnnl::memory::format_tag::acdb;
|
||||
}
|
||||
|
||||
dnnl::memory::data_type xType = dnnl::memory::data_type::f32;
|
||||
|
||||
dnnl::memory::desc x_mkl_md, x_user_md, z_mkl_md, z_user_md;
|
||||
|
||||
x_user_md = x_mkl_md = dnnl::memory::desc(shape, xType, xFormat);
|
||||
onednnUtils::setBlockStrides(*x, x_user_md);
|
||||
|
||||
// z
|
||||
z_user_md = z_mkl_md = dnnl::memory::desc(shape, xType, zFormat);
|
||||
onednnUtils::setBlockStrides(*z, z_user_md);
|
||||
|
||||
auto engine = onednnUtils::getEngine(LaunchContext::defaultContext()->engine());
|
||||
|
||||
// Create attributes (to handle alpha and beta if necessary)
|
||||
dnnl::primitive_attr attr; // it is empty since we have usual values for alpha (=1) and beta (=0)
|
||||
|
||||
// operation primitive description
|
||||
dnnl::softmax_forward::desc op_desc(dnnl::prop_kind::forward_inference, x_mkl_md, axis);
|
||||
|
||||
dnnl::softmax_forward::primitive_desc op_prim_desc(op_desc, attr, engine);
|
||||
|
||||
// arguments (memory buffers) necessary for calculations
|
||||
std::unordered_map<int, dnnl::memory> args;
|
||||
|
||||
dnnl::stream stream(engine);
|
||||
|
||||
// provide memory buffers and check whether reorder is required
|
||||
|
||||
// input
|
||||
onednnUtils::loadDataToMklStream(*x, engine, stream, x_user_md, op_prim_desc.src_desc(), args[DNNL_ARG_SRC]);
|
||||
|
||||
// z
|
||||
auto z_user_mem =
|
||||
onednnUtils::loadDataToMklStream(*z, engine, stream, z_user_md, op_prim_desc.dst_desc(), args[DNNL_ARG_DST]);
|
||||
|
||||
// run calculations
|
||||
dnnl::softmax_forward(op_prim_desc).execute(stream, args);
|
||||
|
||||
// reorder outputs if necessary
|
||||
if (op_prim_desc.dst_desc() != z_user_mem.get_desc())
|
||||
dnnl::reorder(args[DNNL_ARG_DST], z_user_mem).execute(stream, args[DNNL_ARG_DST], z_user_mem);
|
||||
|
||||
stream.wait();
|
||||
}
|
||||
|
||||
PLATFORM_IMPL(softmax, ENGINE_CPU) {
|
||||
auto input = INPUT_VARIABLE(0);
|
||||
auto output = OUTPUT_VARIABLE(0);
|
||||
|
||||
const int rank = input->rankOf();
|
||||
int dim = block.getIArguments()->size() > 0 ? INT_ARG(0) : rank - 1;
|
||||
|
||||
if (dim < 0) {
|
||||
dim += rank;
|
||||
}
|
||||
|
||||
REQUIRE_TRUE(dim < rank && dim >= 0, 0,
|
||||
"SOFTMAX_MKLDNN OP: the value of input integer parameter (dimension) must be less than input array rank "
|
||||
"%i, but got dimension = %i instead !",
|
||||
rank, dim);
|
||||
|
||||
REQUIRE_TRUE(rank <= 6, 0, "SOFTMAX_MKLDNN OP: the rank of input must be less or qual 6, but got rank = %i instead !",
|
||||
rank);
|
||||
|
||||
// mkldnnSoftMax
|
||||
softmaxMKLDNN(input, output, dim);
|
||||
|
||||
return sd::Status::OK;
|
||||
}
|
||||
|
||||
PLATFORM_CHECK(softmax, ENGINE_CPU) {
|
||||
auto x = INPUT_VARIABLE(0);
|
||||
auto z = OUTPUT_VARIABLE(0);
|
||||
|
||||
Requirements req("ONEDNN SOFTMAX OP");
|
||||
req.expectTrue(block.isUseONEDNN(), IS_USE_ONEDNN_MSG) &&
|
||||
req.expectFalse(makeInfoVariable(x->isEmpty(), IS_EMPTY_MSG_INPUT), EXPECTED_FALSE) &&
|
||||
req.expectLess(makeInfoVariable(x->rankOf(), RANK_MSG_INPUT), 7) &&
|
||||
req.expectGreater(makeInfoVariable(x->rankOf(), RANK_MSG_INPUT), 2) &&
|
||||
req.expectEq(makeInfoVariable(x->dataType(), TYPE_MSG_INPUT), DataType::FLOAT32) &&
|
||||
req.expectEq(makeInfoVariable(z->dataType(), TYPE_MSG_OUTPUT), DataType::FLOAT32);
|
||||
req.logTheSuccess();
|
||||
return req;
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////
|
||||
static void softmaxBpMKLDNN(NDArray* x, NDArray* dLdz, NDArray* dLdx, const sd::LongType axis) {
|
||||
dnnl::memory::desc x_user_md, x_mkl_md, dLdx_mkl_md, dLdx_user_md, dLdz_mkl_md, dLdz_user_md;
|
||||
|
||||
// x
|
||||
x_mkl_md = x_user_md =
|
||||
dnnl::memory::desc(x->getShapeAsFlatVector(), dnnl::memory::data_type::f32, onednnUtils::getFormat(*x));
|
||||
onednnUtils::setBlockStrides(*x, x_user_md);
|
||||
|
||||
// dLdx
|
||||
dLdx_mkl_md = dLdx_user_md =
|
||||
dnnl::memory::desc(dLdx->getShapeAsFlatVector(), dnnl::memory::data_type::f32, onednnUtils::getFormat(*dLdx));
|
||||
onednnUtils::setBlockStrides(*dLdx, dLdx_user_md);
|
||||
|
||||
// dLdz
|
||||
dLdz_mkl_md = dLdz_user_md =
|
||||
dnnl::memory::desc(dLdz->getShapeAsFlatVector(), dnnl::memory::data_type::f32, onednnUtils::getFormat(*dLdz));
|
||||
onednnUtils::setBlockStrides(*dLdz, dLdz_user_md);
|
||||
|
||||
auto engine = onednnUtils::getEngine(LaunchContext::defaultContext()->engine());
|
||||
|
||||
// operation primitive description
|
||||
// forward description
|
||||
dnnl::softmax_forward::desc op_ff_desc(dnnl::prop_kind::forward_inference, x_mkl_md, axis);
|
||||
dnnl::softmax_forward::primitive_desc op_ff_prim_desc(op_ff_desc, engine);
|
||||
|
||||
// backward description
|
||||
dnnl::softmax_backward::desc op_bp_desc(dLdz_mkl_md, dLdx_mkl_md, axis);
|
||||
dnnl::softmax_backward::primitive_desc op_bp_prim_desc(op_bp_desc, engine, op_ff_prim_desc);
|
||||
|
||||
// arguments (memory buffers) necessary for calculations
|
||||
std::unordered_map<int, dnnl::memory> argsbp, argsff;
|
||||
|
||||
dnnl::stream stream(engine);
|
||||
|
||||
// provide memory buffers and check whether reorder is required for forward
|
||||
// input
|
||||
onednnUtils::loadDataToMklStream(*x, engine, stream, x_user_md, op_ff_prim_desc.src_desc(), argsff[DNNL_ARG_SRC]);
|
||||
|
||||
// dLdz
|
||||
onednnUtils::loadDataToMklStream(*dLdz, engine, stream, dLdz_user_md, op_bp_prim_desc.diff_dst_desc(),
|
||||
argsbp[DNNL_ARG_DIFF_DST]);
|
||||
|
||||
// dLdx
|
||||
auto dLdx_user_mem = onednnUtils::loadDataToMklStream(*dLdx, engine, stream, dLdx_user_md, op_ff_prim_desc.src_desc(),
|
||||
argsff[DNNL_ARG_DST]);
|
||||
|
||||
// check and arg set for backprob
|
||||
argsbp[DNNL_ARG_DIFF_SRC] = argsff[DNNL_ARG_DST];
|
||||
argsbp[DNNL_ARG_DST] = argsff[DNNL_ARG_DST];
|
||||
|
||||
|
||||
// run calculations backward
|
||||
dnnl::softmax_backward(op_bp_prim_desc).execute(stream, argsbp);
|
||||
|
||||
// reorder outputs if necessary
|
||||
if (op_ff_prim_desc.dst_desc() != dLdx_user_mem.get_desc())
|
||||
dnnl::reorder(argsff[DNNL_ARG_DST], dLdx_user_mem).execute(stream, argsff[DNNL_ARG_DST], dLdx_user_mem);
|
||||
|
||||
stream.wait();
|
||||
}
|
||||
|
||||
PLATFORM_IMPL(softmax_bp, ENGINE_CPU) {
|
||||
auto input = INPUT_VARIABLE(0);
|
||||
auto dLdz = INPUT_VARIABLE(1);
|
||||
auto softmaxOutput = INPUT_VARIABLE(2);
|
||||
auto dLdx = OUTPUT_VARIABLE(0);
|
||||
dLdx->assign(softmaxOutput);
|
||||
|
||||
const sd::LongType rank = input->rankOf();
|
||||
const sd::LongType dLdzRank = dLdz->rankOf();
|
||||
int dim = block.getIArguments()->size() > 0 ? INT_ARG(0) : rank - 1;
|
||||
|
||||
if (dim < 0) {
|
||||
dim += rank;
|
||||
}
|
||||
|
||||
REQUIRE_TRUE(dim < rank && dim >= 0, 0,
|
||||
"SOFTMAX_MKLDNN_BP OP: the value of input integer parameter (dimension) must be less than input array "
|
||||
"rank %i, but got dimension = %i instead !",
|
||||
rank, dim);
|
||||
|
||||
REQUIRE_TRUE(rank <= 6 && dLdzRank <= 6, 0,
|
||||
"SOFTMAX_MKLDNN_BP OP: the rank of input and dLdz must be less or qual 6, but got input rank = %i and "
|
||||
"dLdz rank rank = %i instead !",
|
||||
rank, dLdzRank);
|
||||
|
||||
// mkldnnSoftMax
|
||||
softmaxBpMKLDNN(input, dLdz, dLdx, dim);
|
||||
|
||||
return sd::Status::OK;
|
||||
}
|
||||
|
||||
PLATFORM_CHECK(softmax_bp, ENGINE_CPU) {
|
||||
auto x = INPUT_VARIABLE(0);
|
||||
auto dLdz = INPUT_VARIABLE(1);
|
||||
auto softmaxOutput = INPUT_VARIABLE(2);
|
||||
auto dLdx = OUTPUT_VARIABLE(0);
|
||||
dLdx->assign(softmaxOutput);
|
||||
|
||||
Requirements req("ONEDNN SOFTMAX_BP OP");
|
||||
req.expectTrue(block.isUseONEDNN(), IS_USE_ONEDNN_MSG) &&
|
||||
req.expectFalse(makeInfoVariable(x->isEmpty(), IS_EMPTY_MSG_INPUT0), EXPECTED_FALSE) &&
|
||||
req.expectFalse(makeInfoVariable(dLdz->isEmpty(), IS_EMPTY_MSG_INPUT1), EXPECTED_FALSE) &&
|
||||
req.expectLess(makeInfoVariable(x->rankOf(), RANK_MSG_INPUT0), 7) &&
|
||||
req.expectEq(makeInfoVariable(x->rankOf(), RANK_MSG_INPUT0), makeInfoVariable(dLdz->rankOf(), RANK_MSG_INPUT1)) &&
|
||||
req.expectGreater(makeInfoVariable(x->rankOf(), RANK_MSG_INPUT), 0) &&
|
||||
req.expectEq(makeInfoVariable(x->dataType(), TYPE_MSG_INPUT0), DataType::FLOAT32) &&
|
||||
req.expectEq(makeInfoVariable(dLdz->dataType(), TYPE_MSG_INPUT1), DataType::FLOAT32) &&
|
||||
req.expectEq(makeInfoVariable(dLdx->dataType(), TYPE_MSG_OUTPUT), DataType::FLOAT32) &&
|
||||
req.expect(
|
||||
makeShapeInfoVariable(x, SHAPE_MSG_INPUT0), makeShapeInfoVariable(dLdz, SHAPE_MSG_INPUT1),
|
||||
[](const decltype(x)& l, const decltype(dLdz)& r) {
|
||||
for (int i = 0; i < l->rankOf(); i++) {
|
||||
if (l->sizeAt(i) != r->sizeAt(i)) {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
return true;
|
||||
},
|
||||
EXPECTED_EQ_MSG);
|
||||
req.logTheSuccess();
|
||||
return req;
|
||||
}
|
||||
|
||||
} // namespace platforms
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,218 @@
|
||||
/*
|
||||
* ******************************************************************************
|
||||
* *
|
||||
* *
|
||||
* * 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 Oleg Semeniv <oleg.semeniv@gmail.com>
|
||||
//
|
||||
//
|
||||
#include <helpers/MKLDNNStream.h>
|
||||
#include <ops/declarable/OpRegistrator.h>
|
||||
#include <ops/declarable/PlatformHelper.h>
|
||||
#include <system/platform_boilerplate.h>
|
||||
|
||||
#include "mkldnnUtils.h"
|
||||
|
||||
using namespace dnnl;
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace platforms {
|
||||
|
||||
//////////////////////////////////////////////////////////////////////
|
||||
static void tanhMKLDNN(NDArray* x, NDArray* z) {
|
||||
dnnl::memory::dims shape = x->getShapeAsFlatVector();
|
||||
|
||||
dnnl::memory::desc x_mkl_md, x_user_md, z_mkl_md, z_user_md;
|
||||
|
||||
x_user_md = x_mkl_md = dnnl::memory::desc(shape, dnnl::memory::data_type::f32, onednnUtils::getFormat(*x));
|
||||
onednnUtils::setBlockStrides(*x, x_user_md);
|
||||
|
||||
// z
|
||||
z_user_md = z_mkl_md = dnnl::memory::desc(shape, dnnl::memory::data_type::f32, onednnUtils::getFormat(*z));
|
||||
onednnUtils::setBlockStrides(*z, z_user_md);
|
||||
|
||||
auto engine = onednnUtils::getEngine(LaunchContext::defaultContext()->engine());
|
||||
|
||||
// Create attributes (to handle alpha and beta if necessary)
|
||||
dnnl::primitive_attr attr; // it is empty since we have usual values for alpha (=1) and beta (=0)
|
||||
|
||||
// operation primitive description
|
||||
dnnl::eltwise_forward::desc op_desc(dnnl::prop_kind::forward_inference, algorithm::eltwise_tanh, x_mkl_md, 0, 0);
|
||||
|
||||
dnnl::eltwise_forward::primitive_desc op_prim_desc(op_desc, attr, engine);
|
||||
|
||||
// arguments (memory buffers) necessary for calculations
|
||||
std::unordered_map<int, dnnl::memory> args;
|
||||
|
||||
dnnl::stream stream(engine);
|
||||
|
||||
// provide memory buffers and check whether reorder is required
|
||||
// input
|
||||
onednnUtils::loadDataToMklStream(*x, engine, stream, x_user_md, op_prim_desc.src_desc(), args[DNNL_ARG_SRC]);
|
||||
|
||||
// z
|
||||
auto z_user_mem =
|
||||
onednnUtils::loadDataToMklStream(*z, engine, stream, z_user_md, op_prim_desc.dst_desc(), args[DNNL_ARG_DST]);
|
||||
|
||||
// run calculations
|
||||
dnnl::eltwise_forward(op_prim_desc).execute(stream, args);
|
||||
|
||||
// reorder outputs if necessary
|
||||
if (op_prim_desc.dst_desc() != z_user_mem.get_desc())
|
||||
dnnl::reorder(args[DNNL_ARG_DST], z_user_mem).execute(stream, args[DNNL_ARG_DST], z_user_mem);
|
||||
|
||||
stream.wait();
|
||||
}
|
||||
|
||||
PLATFORM_IMPL(tanh, ENGINE_CPU) {
|
||||
auto input = INPUT_VARIABLE(0);
|
||||
auto output = OUTPUT_VARIABLE(0);
|
||||
const sd::LongType rank = input->rankOf();
|
||||
REQUIRE_TRUE(rank <= 6, 0, "TANH_MKLDNN OP: the rank of input must be less or qual 6, but got rank = %i instead !",
|
||||
rank);
|
||||
|
||||
// mkldnnTanh
|
||||
tanhMKLDNN(input, output);
|
||||
|
||||
return sd::Status::OK;
|
||||
}
|
||||
|
||||
PLATFORM_CHECK(tanh, ENGINE_CPU) {
|
||||
auto x = INPUT_VARIABLE(0);
|
||||
auto z = OUTPUT_VARIABLE(0);
|
||||
|
||||
Requirements req("ONEDNN TANH OP");
|
||||
req.expectTrue(block.isUseONEDNN(), IS_USE_ONEDNN_MSG) &&
|
||||
req.expectFalse(makeInfoVariable(x->isEmpty(), IS_EMPTY_MSG_INPUT), EXPECTED_FALSE) &&
|
||||
req.expectLess(makeInfoVariable(x->rankOf(), RANK_MSG_INPUT), 7) &&
|
||||
req.expectGreater(makeInfoVariable(x->rankOf(), RANK_MSG_INPUT), 0) &&
|
||||
req.expectEq(makeInfoVariable(x->dataType(), TYPE_MSG_INPUT), DataType::FLOAT32) &&
|
||||
req.expectEq(makeInfoVariable(z->dataType(), TYPE_MSG_OUTPUT), DataType::FLOAT32);
|
||||
req.logTheSuccess();
|
||||
return req;
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////
|
||||
static void tanhBpMKLDNN(NDArray* x, NDArray* dLdz, NDArray* dLdx) {
|
||||
dnnl::memory::dims shape = x->getShapeAsFlatVector();
|
||||
|
||||
dnnl::memory::desc x_mkl_md, x_user_md, dLdx_mkl_md, dLdx_user_md, dLdz_mkl_md, dLdz_user_md;
|
||||
|
||||
// x
|
||||
x_user_md = x_mkl_md = dnnl::memory::desc(shape, dnnl::memory::data_type::f32, onednnUtils::getFormat(*x));
|
||||
onednnUtils::setBlockStrides(*x, x_user_md);
|
||||
|
||||
// dLdz
|
||||
dLdz_user_md = dLdz_mkl_md = dnnl::memory::desc(shape, dnnl::memory::data_type::f32, onednnUtils::getFormat(*dLdz));
|
||||
onednnUtils::setBlockStrides(*dLdz, dLdz_user_md);
|
||||
|
||||
// dLdx
|
||||
dLdx_user_md = dLdx_mkl_md = dnnl::memory::desc(shape, dnnl::memory::data_type::f32, onednnUtils::getFormat(*dLdx));
|
||||
onednnUtils::setBlockStrides(*dLdx, dLdx_user_md);
|
||||
|
||||
auto engine = onednnUtils::getEngine(LaunchContext::defaultContext()->engine());
|
||||
|
||||
// arguments (memory buffers) necessary for calculations
|
||||
std::unordered_map<int, dnnl::memory> args;
|
||||
|
||||
dnnl::stream stream(engine);
|
||||
|
||||
// operation primitive description
|
||||
// forward
|
||||
dnnl::eltwise_forward::desc op_ff_desc(dnnl::prop_kind::forward_inference, algorithm::eltwise_tanh, x_mkl_md, 0, 0);
|
||||
dnnl::eltwise_forward::primitive_desc op_ff_prim_desc(op_ff_desc, engine);
|
||||
|
||||
// backward description
|
||||
dnnl::eltwise_backward::desc op_desc(algorithm::eltwise_tanh, dLdz_mkl_md, x_mkl_md, 0, 0);
|
||||
dnnl::eltwise_backward::primitive_desc op_prim_desc(op_desc, engine, op_ff_prim_desc);
|
||||
|
||||
// provide memory buffers and check whether reorder is required for forward
|
||||
// input
|
||||
onednnUtils::loadDataToMklStream(*x, engine, stream, x_user_md, op_prim_desc.src_desc(), args[DNNL_ARG_SRC]);
|
||||
|
||||
// dLdz
|
||||
onednnUtils::loadDataToMklStream(*dLdz, engine, stream, dLdz_user_md, op_prim_desc.diff_dst_desc(),
|
||||
args[DNNL_ARG_DIFF_DST]);
|
||||
|
||||
// dLdx
|
||||
auto dLdx_user_mem = onednnUtils::loadDataToMklStream(*dLdx, engine, stream, dLdx_user_md,
|
||||
op_prim_desc.diff_src_desc(), args[DNNL_ARG_DIFF_SRC]);
|
||||
|
||||
// run calculations backward
|
||||
dnnl::eltwise_backward(op_prim_desc).execute(stream, args);
|
||||
|
||||
// reorder outputs if necessary
|
||||
if (op_prim_desc.diff_src_desc() != dLdx_user_mem.get_desc())
|
||||
dnnl::reorder(args[DNNL_ARG_DIFF_SRC], dLdx_user_mem).execute(stream, args[DNNL_ARG_DIFF_SRC], dLdx_user_mem);
|
||||
|
||||
stream.wait();
|
||||
}
|
||||
|
||||
PLATFORM_IMPL(tanh_bp, ENGINE_CPU) {
|
||||
auto input = INPUT_VARIABLE(0);
|
||||
auto dLdz = INPUT_VARIABLE(1);
|
||||
auto dLdx = OUTPUT_VARIABLE(0);
|
||||
|
||||
const sd::LongType rank = input->rankOf();
|
||||
const sd::LongType dLdzRank = dLdz->rankOf();
|
||||
|
||||
REQUIRE_TRUE(rank <= 6 && dLdzRank <= 6, 0,
|
||||
"TANH_BP_MKLDNN OP: the rank of input and dLdz must be less or qual 6, but got input rank = %i and dLdz "
|
||||
"rank rank = %i instead !",
|
||||
rank, dLdzRank);
|
||||
|
||||
// mkldnnSoftMax
|
||||
tanhBpMKLDNN(input, dLdz, dLdx);
|
||||
|
||||
return sd::Status::OK;
|
||||
}
|
||||
|
||||
PLATFORM_CHECK(tanh_bp, ENGINE_CPU) {
|
||||
auto x = INPUT_VARIABLE(0);
|
||||
auto dLdz = INPUT_VARIABLE(1);
|
||||
auto dLdx = OUTPUT_VARIABLE(0);
|
||||
|
||||
Requirements req("ONEDNN TANH BP OP");
|
||||
req.expectTrue(block.isUseONEDNN(), IS_USE_ONEDNN_MSG) &&
|
||||
req.expectFalse(makeInfoVariable(x->isEmpty(), IS_EMPTY_MSG_INPUT0), EXPECTED_FALSE) &&
|
||||
req.expectFalse(makeInfoVariable(dLdz->isEmpty(), IS_EMPTY_MSG_INPUT1), EXPECTED_FALSE) &&
|
||||
req.expectLess(makeInfoVariable(x->rankOf(), RANK_MSG_INPUT0), 7) &&
|
||||
req.expectGreater(makeInfoVariable(x->rankOf(), RANK_MSG_INPUT0), 0) &&
|
||||
req.expectEq(makeInfoVariable(x->dataType(), TYPE_MSG_INPUT0), DataType::FLOAT32) &&
|
||||
req.expectEq(makeInfoVariable(dLdz->dataType(), TYPE_MSG_INPUT1), DataType::FLOAT32) &&
|
||||
req.expectEq(makeInfoVariable(dLdx->dataType(), TYPE_MSG_OUTPUT), DataType::FLOAT32) &&
|
||||
req.expect(
|
||||
makeShapeInfoVariable(x, SHAPE_MSG_INPUT0), makeShapeInfoVariable(dLdz, SHAPE_MSG_INPUT1),
|
||||
[](const decltype(x)& l, const decltype(dLdz)& r) {
|
||||
for (int i = 0; i < l->rankOf(); i++) {
|
||||
if (l->sizeAt(i) != r->sizeAt(i)) {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
return true;
|
||||
},
|
||||
EXPECTED_EQ_MSG);
|
||||
req.logTheSuccess();
|
||||
return req;
|
||||
}
|
||||
|
||||
} // namespace platforms
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
@@ -0,0 +1,415 @@
|
||||
/*
|
||||
* ******************************************************************************
|
||||
* *
|
||||
* *
|
||||
* * 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 Oleg Semeniv <oleg.semeniv@gmail.com>
|
||||
//
|
||||
//
|
||||
#include <helpers/MKLDNNStream.h>
|
||||
#include <ops/declarable/OpRegistrator.h>
|
||||
#include <ops/declarable/PlatformHelper.h>
|
||||
#include <system/platform_boilerplate.h>
|
||||
|
||||
#include "mkldnnUtils.h"
|
||||
|
||||
using namespace dnnl;
|
||||
|
||||
namespace sd {
|
||||
namespace ops {
|
||||
namespace platforms {
|
||||
|
||||
//////////////////////////////////////////////////////////////////////
|
||||
static void xwPlusBiasMKLDNN(NDArray* x, NDArray* weights, NDArray* bias, NDArray* z,
|
||||
const bool bShouldTransp) {
|
||||
// mkl works with following
|
||||
// [M,K] x [N,K]^T + [N] = [M,N]
|
||||
const auto xRank = x->rankOf();
|
||||
|
||||
// [M,K] x [K,N] = [M,N]
|
||||
const sd::LongType M = x->sizeAt(0);
|
||||
const sd::LongType K = x->sizeAt(1); // K == wK
|
||||
const sd::LongType N = z->sizeAt(1);
|
||||
|
||||
dnnl::memory::dims xShape = dnnl::memory::dims({M, K});
|
||||
dnnl::memory::dims wShape = dnnl::memory::dims({N, K});
|
||||
dnnl::memory::dims zShape = dnnl::memory::dims({M, N});
|
||||
dnnl::memory::dims bShape = dnnl::memory::dims({N});
|
||||
|
||||
dnnl::memory::format_tag format = dnnl::memory::format_tag::ab;
|
||||
|
||||
// x type
|
||||
dnnl::memory::data_type xType = dnnl::memory::data_type::f32;
|
||||
if (x->dataType() == DataType::UINT8)
|
||||
xType = dnnl::memory::data_type::u8;
|
||||
else if (x->dataType() == DataType::INT8)
|
||||
xType = dnnl::memory::data_type::s8;
|
||||
|
||||
// weights type
|
||||
dnnl::memory::data_type wType = (weights->dataType() == DataType::FLOAT32) ? wType = dnnl::memory::data_type::f32
|
||||
: wType = dnnl::memory::data_type::s8;
|
||||
|
||||
// bias type need add description for bias
|
||||
dnnl::memory::data_type bType = dnnl::memory::data_type::f32;
|
||||
if (bias->dataType() == DataType::INT32)
|
||||
bType = dnnl::memory::data_type::s32;
|
||||
else if (bias->dataType() == DataType::UINT8)
|
||||
bType = dnnl::memory::data_type::u8;
|
||||
else if (bias->dataType() == DataType::INT8)
|
||||
bType = dnnl::memory::data_type::s8;
|
||||
|
||||
// z type
|
||||
dnnl::memory::data_type zType = dnnl::memory::data_type::f32;
|
||||
if (z->dataType() == DataType::INT32)
|
||||
zType = dnnl::memory::data_type::s32;
|
||||
else if (z->dataType() == DataType::UINT8)
|
||||
zType = dnnl::memory::data_type::u8;
|
||||
else if (z->dataType() == DataType::INT8)
|
||||
zType = dnnl::memory::data_type::s8;
|
||||
|
||||
// memory descriptors for arrays
|
||||
// x
|
||||
dnnl::memory::desc x_mkl_md = dnnl::memory::desc(xShape, xType, dnnl::memory::format_tag::any);
|
||||
dnnl::memory::desc x_user_md = dnnl::memory::desc(xShape, xType, onednnUtils::getFormat(*x));
|
||||
onednnUtils::setBlockStrides(*x, x_user_md);
|
||||
|
||||
// weights
|
||||
dnnl::memory::desc weights_mkl_md = dnnl::memory::desc(wShape, wType, dnnl::memory::format_tag::any);
|
||||
dnnl::memory::desc weights_user_md = dnnl::memory::desc(wShape, wType, onednnUtils::getFormat(*weights));
|
||||
onednnUtils::setBlockStrides(*weights, weights_user_md,
|
||||
bShouldTransp ? std::vector<int>({1, 0}) : std::vector<int>());
|
||||
|
||||
// bias
|
||||
dnnl::memory::desc bias_mkl_md = dnnl::memory::desc(bShape, bType, dnnl::memory::format_tag::a);
|
||||
dnnl::memory::desc bias_user_md = dnnl::memory::desc(bShape, bType, dnnl::memory::format_tag::a);
|
||||
onednnUtils::setBlockStrides(*bias, bias_user_md);
|
||||
|
||||
// z
|
||||
dnnl::memory::desc z_mkl_md = dnnl::memory::desc(zShape, zType, dnnl::memory::format_tag::any);
|
||||
dnnl::memory::desc z_user_md = dnnl::memory::desc(zShape, zType, onednnUtils::getFormat(*z));
|
||||
onednnUtils::setBlockStrides(*z, z_user_md);
|
||||
|
||||
auto engine = onednnUtils::getEngine(LaunchContext::defaultContext()->engine());
|
||||
|
||||
// operation primitive description
|
||||
dnnl::inner_product_forward::desc op_desc(dnnl::prop_kind::forward_inference, x_mkl_md, weights_mkl_md, bias_mkl_md,
|
||||
z_mkl_md);
|
||||
|
||||
dnnl::inner_product_forward::primitive_desc op_prim_desc(op_desc, engine);
|
||||
|
||||
// arguments (memory buffers) necessary for calculations
|
||||
std::unordered_map<int, dnnl::memory> args;
|
||||
|
||||
dnnl::stream stream(engine);
|
||||
|
||||
// provide memory buffers and check whether reorder is required
|
||||
|
||||
// input
|
||||
onednnUtils::loadDataToMklStream(*x, engine, stream, x_user_md, op_prim_desc.src_desc(), args[DNNL_ARG_SRC]);
|
||||
|
||||
// weights
|
||||
onednnUtils::loadDataToMklStream(*weights, engine, stream, weights_user_md, op_prim_desc.weights_desc(),
|
||||
args[DNNL_ARG_WEIGHTS]);
|
||||
|
||||
// bias
|
||||
auto bias_mkl_mem = dnnl::memory(bias_mkl_md, engine, const_cast<void*>(bias->buffer()));
|
||||
args[DNNL_ARG_BIAS] = bias_mkl_mem;
|
||||
|
||||
// z
|
||||
auto z_user_mem =
|
||||
onednnUtils::loadDataToMklStream(*z, engine, stream, z_user_md, op_prim_desc.dst_desc(), args[DNNL_ARG_DST]);
|
||||
|
||||
// run calculations
|
||||
dnnl::inner_product_forward(op_prim_desc).execute(stream, args);
|
||||
|
||||
// reorder outputs if necessary
|
||||
if (op_prim_desc.dst_desc() != z_user_mem.get_desc())
|
||||
dnnl::reorder(args[DNNL_ARG_DST], z_user_mem).execute(stream, args[DNNL_ARG_DST], z_user_mem);
|
||||
|
||||
stream.wait();
|
||||
}
|
||||
|
||||
//////////////////////////////////////////////////////////////////////
|
||||
static void xwPlusBiasBp(NDArray* x, NDArray* weights, NDArray* bias, NDArray* dLdz,
|
||||
NDArray* dLdx, NDArray* dLdw, NDArray* dLdb, const bool bShouldTransp) {
|
||||
// mkl works with following
|
||||
// [M,K] x [N,K]^T + [N] = [M,N]
|
||||
const auto xRank = x->rankOf();
|
||||
|
||||
// [M,K] x [K,N] = [M,N]
|
||||
const sd::LongType M = x->sizeAt(0);
|
||||
const sd::LongType K = x->sizeAt(1); // K == wK
|
||||
const sd::LongType N = dLdz->sizeAt(1);
|
||||
// input dims
|
||||
dnnl::memory::dims xShape = dnnl::memory::dims({M, K});
|
||||
dnnl::memory::dims wShape = dnnl::memory::dims({N, K});
|
||||
dnnl::memory::dims dLdzShape = dnnl::memory::dims({M, N});
|
||||
|
||||
dnnl::memory::dims bShape = dnnl::memory::dims({N});
|
||||
|
||||
// output dims
|
||||
dnnl::memory::dims dLdxShape = xShape;
|
||||
dnnl::memory::dims dLdwShape = wShape;
|
||||
|
||||
dnnl::memory::data_type dataType = dnnl::memory::data_type::f32;
|
||||
|
||||
// memory descriptors for arrays
|
||||
// x
|
||||
dnnl::memory::desc x_mkl_md = dnnl::memory::desc(xShape, dataType, dnnl::memory::format_tag::any);
|
||||
dnnl::memory::desc x_user_md = dnnl::memory::desc(xShape, dataType, onednnUtils::getFormat(*x));
|
||||
onednnUtils::setBlockStrides(*x, x_user_md);
|
||||
|
||||
// weights
|
||||
dnnl::memory::desc weights_mkl_md = dnnl::memory::desc(wShape, dataType, dnnl::memory::format_tag::any);
|
||||
dnnl::memory::desc weights_user_md = dnnl::memory::desc(wShape, dataType, onednnUtils::getFormat(*weights));
|
||||
onednnUtils::setBlockStrides(*weights, weights_user_md,
|
||||
bShouldTransp ? std::vector<int>({1, 0}) : std::vector<int>());
|
||||
|
||||
// bias
|
||||
dnnl::memory::desc bias_mkl_md = dnnl::memory::desc(bShape, dataType, dnnl::memory::format_tag::any);
|
||||
dnnl::memory::desc bias_user_md = dnnl::memory::desc(bShape, dataType, onednnUtils::getFormat(*bias));
|
||||
onednnUtils::setBlockStrides(*bias, bias_user_md);
|
||||
|
||||
// dLdz
|
||||
dnnl::memory::desc dLdz_mkl_md = dnnl::memory::desc(dLdzShape, dataType, dnnl::memory::format_tag::any);
|
||||
dnnl::memory::desc dLdz_user_md = dnnl::memory::desc(dLdzShape, dataType, onednnUtils::getFormat(*dLdz));
|
||||
onednnUtils::setBlockStrides(*dLdz, dLdz_user_md);
|
||||
|
||||
// dLdw
|
||||
dnnl::memory::desc dLdw_mkl_md = dnnl::memory::desc(wShape, dataType, dnnl::memory::format_tag::any);
|
||||
dnnl::memory::desc dLdw_user_md = dnnl::memory::desc(wShape, dataType, onednnUtils::getFormat(*dLdw));
|
||||
onednnUtils::setBlockStrides(*dLdw, dLdw_user_md, bShouldTransp ? std::vector<int>({1, 0}) : std::vector<int>());
|
||||
|
||||
// dLdb
|
||||
dnnl::memory::desc dLdb_mkl_md = dnnl::memory::desc(bShape, dataType, dnnl::memory::format_tag::any);
|
||||
dnnl::memory::desc dLdb_user_md = dnnl::memory::desc(bShape, dataType, onednnUtils::getFormat(*dLdb));
|
||||
onednnUtils::setBlockStrides(*dLdb, dLdb_user_md);
|
||||
|
||||
// dLdx
|
||||
dnnl::memory::desc dLdx_mkl_md = dnnl::memory::desc(xShape, dataType, dnnl::memory::format_tag::any);
|
||||
dnnl::memory::desc dLdx_user_md = dnnl::memory::desc(xShape, dataType, onednnUtils::getFormat(*dLdx));
|
||||
onednnUtils::setBlockStrides(*dLdx, dLdx_user_md);
|
||||
|
||||
// create engine
|
||||
auto engine = onednnUtils::getEngine(LaunchContext::defaultContext()->engine());
|
||||
|
||||
// forward
|
||||
// operation primitive description
|
||||
dnnl::inner_product_forward::desc op_ff_desc(dnnl::prop_kind::forward_inference, x_mkl_md, weights_mkl_md,
|
||||
bias_mkl_md, dLdz_mkl_md);
|
||||
dnnl::inner_product_forward::primitive_desc op_ff_prim_desc(op_ff_desc, engine);
|
||||
|
||||
// backprob
|
||||
// dLdw
|
||||
auto op_bpdw_desc = inner_product_backward_weights::desc(x_mkl_md, dLdw_mkl_md, dLdb_mkl_md, dLdz_mkl_md);
|
||||
auto op_bpdw_prim_desc = inner_product_backward_weights::primitive_desc(op_bpdw_desc, engine, op_ff_prim_desc);
|
||||
|
||||
// backprob
|
||||
// dLdx
|
||||
auto op_bpdx_desc = inner_product_backward_data::desc(dLdx_mkl_md, weights_mkl_md, dLdz_mkl_md);
|
||||
auto op_bpdx_prim_desc = inner_product_backward_data::primitive_desc(op_bpdx_desc, engine, op_ff_prim_desc);
|
||||
|
||||
// arguments (memory buffers) necessary for calculations
|
||||
std::unordered_map<int, dnnl::memory> argsDw, argsDx;
|
||||
|
||||
dnnl::stream stream(engine);
|
||||
|
||||
// dLdz dw
|
||||
onednnUtils::loadDataToMklStream(*dLdz, engine, stream, dLdz_user_md, op_bpdw_prim_desc.diff_dst_desc(),
|
||||
argsDw[DNNL_ARG_DIFF_DST]);
|
||||
|
||||
// dLdz - dx
|
||||
onednnUtils::loadDataToMklStream(*dLdz, engine, stream, dLdz_user_md, op_bpdx_prim_desc.diff_dst_desc(),
|
||||
argsDx[DNNL_ARG_DIFF_DST]);
|
||||
|
||||
// input x for dw
|
||||
onednnUtils::loadDataToMklStream(*x, engine, stream, x_user_md, op_bpdw_prim_desc.src_desc(), argsDw[DNNL_ARG_SRC]);
|
||||
|
||||
// weights - dx
|
||||
onednnUtils::loadDataToMklStream(*weights, engine, stream, weights_user_md, op_bpdx_prim_desc.weights_desc(),
|
||||
argsDx[DNNL_ARG_WEIGHTS]);
|
||||
|
||||
// dLdw
|
||||
auto dLdw_user_mem = onednnUtils::loadDataToMklStream(
|
||||
*dLdw, engine, stream, dLdw_user_md, op_bpdw_prim_desc.diff_weights_desc(), argsDw[DNNL_ARG_DIFF_WEIGHTS]);
|
||||
|
||||
// dLdx
|
||||
auto dLdx_user_mem = onednnUtils::loadDataToMklStream(*dLdx, engine, stream, dLdx_user_md,
|
||||
op_bpdx_prim_desc.diff_src_desc(), argsDx[DNNL_ARG_DIFF_SRC]);
|
||||
|
||||
// dLdb
|
||||
auto dLdb_user_mem = onednnUtils::loadDataToMklStream(*dLdb, engine, stream, dLdb_user_md,
|
||||
op_bpdw_prim_desc.diff_bias_desc(), argsDw[DNNL_ARG_DIFF_BIAS]);
|
||||
|
||||
// run calculations dw
|
||||
dnnl::inner_product_backward_weights(op_bpdw_prim_desc).execute(stream, argsDw);
|
||||
// run calculations dx
|
||||
dnnl::inner_product_backward_data(op_bpdx_prim_desc).execute(stream, argsDx);
|
||||
|
||||
// reorder outputs if necessary
|
||||
if (op_bpdx_prim_desc.diff_src_desc() != dLdx_user_mem.get_desc())
|
||||
dnnl::reorder(argsDx[DNNL_ARG_DIFF_SRC], dLdx_user_mem).execute(stream, argsDx[DNNL_ARG_DIFF_SRC], dLdx_user_mem);
|
||||
|
||||
if (op_bpdw_prim_desc.diff_weights_desc() != dLdw_user_mem.get_desc())
|
||||
dnnl::reorder(argsDw[DNNL_ARG_DIFF_WEIGHTS], dLdw_user_mem)
|
||||
.execute(stream, argsDw[DNNL_ARG_DIFF_WEIGHTS], dLdw_user_mem);
|
||||
|
||||
if (op_bpdw_prim_desc.diff_bias_desc() != dLdb_user_mem.get_desc())
|
||||
dnnl::reorder(argsDw[DNNL_ARG_DIFF_BIAS], dLdb_user_mem).execute(stream, argsDw[DNNL_ARG_DIFF_BIAS], dLdb_user_mem);
|
||||
|
||||
stream.wait();
|
||||
}
|
||||
|
||||
PLATFORM_IMPL(xw_plus_b, ENGINE_CPU) {
|
||||
auto x = INPUT_VARIABLE(0);
|
||||
auto w = INPUT_VARIABLE(1);
|
||||
auto b = INPUT_VARIABLE(2);
|
||||
auto z = OUTPUT_VARIABLE(0);
|
||||
|
||||
if (x->isEmpty() || w->isEmpty() || b->isEmpty()) return sd::Status::OK;
|
||||
|
||||
const int xRank = x->rankOf();
|
||||
const int wRank = w->rankOf();
|
||||
const int zRank = z->rankOf();
|
||||
|
||||
const bool bShouldTransp = block.getIArguments()->size() > 0
|
||||
? (1 != INT_ARG(0))
|
||||
: true; // [M,K] * [K,N] -> [M, N], mkl -> [M,K] * [N, K]^T -> [M, N]
|
||||
|
||||
REQUIRE_TRUE(xRank == 2, 0, "xw_plus_b MKL: Input x array should have rank equal 2, but got instead %i!", xRank);
|
||||
REQUIRE_TRUE(wRank == 2, 0, "xw_plus_b MKL: Input weights array should have rank equal 2, but got instead %i!",
|
||||
wRank);
|
||||
REQUIRE_TRUE(zRank == 2, 0, "xw_plus_b MKL: Output array should have rank equal 2, but got instead %i!", zRank);
|
||||
|
||||
REQUIRE_TRUE(1 == b->rankOf() && b->lengthOf() == z->sizeAt(-1), 0,
|
||||
"xw_plus_b MKL: Input bias vector should be 1D and have proper dimension 1x%i."
|
||||
" But got rank %i, and got length %i instead %i.",
|
||||
z->sizeAt(-1), b->rankOf(), b->lengthOf(), z->sizeAt(-1));
|
||||
|
||||
// mkldnnInerPorductss
|
||||
xwPlusBiasMKLDNN(x, w, b, z, bShouldTransp);
|
||||
|
||||
return sd::Status::OK;
|
||||
}
|
||||
|
||||
PLATFORM_CHECK(xw_plus_b, ENGINE_CPU) {
|
||||
auto x = INPUT_VARIABLE(0);
|
||||
auto w = INPUT_VARIABLE(1);
|
||||
auto b = INPUT_VARIABLE(2);
|
||||
auto z = OUTPUT_VARIABLE(0);
|
||||
|
||||
Requirements req("ONEDNN XW_PLUS_B OP");
|
||||
req.expectTrue(block.isUseONEDNN(), IS_USE_ONEDNN_MSG) &&
|
||||
req.expectTrue(makeInfoVariable(
|
||||
[x, w, b, z] {
|
||||
const DataType xType = x->dataType();
|
||||
const DataType wType = w->dataType();
|
||||
const DataType bType = b->dataType();
|
||||
const DataType zType = z->dataType();
|
||||
|
||||
return ((xType == DataType::FLOAT32 && wType == DataType::FLOAT32 &&
|
||||
bType == DataType::FLOAT32 && zType == DataType::FLOAT32) ||
|
||||
( // x
|
||||
(xType == DataType::UINT8 || xType == DataType::INT8) &&
|
||||
// w
|
||||
(wType == DataType::UINT8 || wType == DataType::INT8) &&
|
||||
// b
|
||||
(bType == DataType::UINT8 || bType == DataType::INT8 ||
|
||||
bType == DataType::INT32 || bType == DataType::FLOAT32) &&
|
||||
// z
|
||||
(zType == DataType::UINT8 || zType == DataType::INT8 ||
|
||||
zType == DataType::INT32 || zType == DataType::FLOAT32)));
|
||||
},
|
||||
TYPECHECK_MSG),
|
||||
NO_MSG);
|
||||
req.logTheSuccess();
|
||||
return req;
|
||||
}
|
||||
|
||||
PLATFORM_IMPL(xw_plus_b_bp, ENGINE_CPU) {
|
||||
auto x = INPUT_VARIABLE(0);
|
||||
auto w = INPUT_VARIABLE(1);
|
||||
auto b = INPUT_VARIABLE(2);
|
||||
auto dLdz = INPUT_VARIABLE(3);
|
||||
|
||||
auto dLdx = OUTPUT_VARIABLE(0);
|
||||
auto dLdw = OUTPUT_VARIABLE(1);
|
||||
auto dLdb = OUTPUT_VARIABLE(2);
|
||||
|
||||
if (x->isEmpty() || w->isEmpty() || b->isEmpty() || dLdz->isEmpty()) return sd::Status::OK;
|
||||
|
||||
const int xRank = x->rankOf();
|
||||
const int wRank = w->rankOf();
|
||||
const int dLdzRank = dLdz->rankOf();
|
||||
|
||||
const bool bShouldTransp = block.getIArguments()->size() > 0
|
||||
? (1 != INT_ARG(0))
|
||||
: true; // [M,K] * [K,N] -> [M, N], mkl -> [M,K] * [N, K]^T -> [M, N]
|
||||
|
||||
REQUIRE_TRUE(x->rankOf() == 2, 0, "xw_plus_b BP MKL: Input x array should have rank equal 2, but got instead %i!",
|
||||
x->rankOf());
|
||||
REQUIRE_TRUE(w->rankOf() == 2, 0,
|
||||
"xw_plus_b BP MKL: Input weights array should have rank equal 2, but got instead %i!", w->rankOf());
|
||||
REQUIRE_TRUE(dLdz->rankOf() == 2, 0, "xw_plus_b BP MKL: Output array should have rank equal 2, but got instead %i!",
|
||||
dLdz->rankOf());
|
||||
REQUIRE_TRUE(1 == b->rankOf() && b->lengthOf() == dLdz->sizeAt(1), 0,
|
||||
"xw_plus_b BP MKL: Input bias vector should be 1D and have proper dimension 1x%i."
|
||||
" But got rank %i, and got length %i instead %i.",
|
||||
dLdz->sizeAt(1), b->rankOf(), b->lengthOf(), dLdz->sizeAt(1));
|
||||
|
||||
xwPlusBiasBp(x, w, b, dLdz, dLdx, dLdw, dLdb, bShouldTransp);
|
||||
|
||||
return sd::Status::OK;
|
||||
}
|
||||
|
||||
PLATFORM_CHECK(xw_plus_b_bp, ENGINE_CPU) {
|
||||
auto x = INPUT_VARIABLE(0);
|
||||
auto w = INPUT_VARIABLE(1);
|
||||
auto b = INPUT_VARIABLE(2);
|
||||
auto dLdz = INPUT_VARIABLE(3);
|
||||
|
||||
auto dLdx = OUTPUT_VARIABLE(0);
|
||||
auto dLdw = OUTPUT_VARIABLE(1);
|
||||
auto dLdb = OUTPUT_VARIABLE(2);
|
||||
|
||||
Requirements req("ONEDNN XW_PLUS_B_BP OP");
|
||||
req.expectTrue(block.isUseONEDNN(), IS_USE_ONEDNN_MSG) &&
|
||||
req.expectTrue(makeInfoVariable(
|
||||
[x, w, b, dLdz, dLdx, dLdw, dLdb] {
|
||||
const DataType xType = x->dataType();
|
||||
const DataType wType = w->dataType();
|
||||
const DataType bType = b->dataType();
|
||||
const DataType dLdzType = dLdz->dataType();
|
||||
const DataType dLdxType = dLdx->dataType();
|
||||
const DataType dLdwType = dLdw->dataType();
|
||||
const DataType dLdbType = dLdb->dataType();
|
||||
return (xType == DataType::FLOAT32 && wType == DataType::FLOAT32 &&
|
||||
bType == DataType::FLOAT32 && dLdzType == DataType::FLOAT32 &&
|
||||
dLdbType == DataType::FLOAT32 && dLdxType == DataType::FLOAT32 &&
|
||||
dLdwType == DataType::FLOAT32);
|
||||
},
|
||||
TYPECHECK_MSG),
|
||||
NO_MSG);
|
||||
req.logTheSuccess();
|
||||
return req;
|
||||
}
|
||||
|
||||
} // namespace platforms
|
||||
} // namespace ops
|
||||
} // namespace sd
|
||||
Reference in New Issue
Block a user