wehub/deeplearning4j--deeplearning4j
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 12:47:05 +08:00
commit 4f3b7da785
7394 changed files with 2005594 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libnd4j/include/ops/declarable/platform/README.md
+1
View File
@@ -0,0 +1 @@
This folder contains platform-specific optimized implementations for custom operations
libnd4j/include/ops/declarable/platform/armcompute/armcomputeUtils.cpp
+261
View File
@@ -0,0 +1,261 @@
/*
* ******************************************************************************
* *
* *
* * 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
* *****************************************************************************
*/
// Created by Abdelrauf 2020
#include "armcomputeUtils.h"
#include <helpers/LoopsCoordsHelper.h>
#include <ops/declarable/OpRegistrator.h>
#include <ops/declarable/PlatformHelper.h>
#include <ops/declarable/helpers/convolutions.h>
#include <system/platform_boilerplate.h>
#include <cstdint>
namespace sd {
namespace ops {
namespace platforms {
Arm_DataType getArmType(const DataType& dType) {
Arm_DataType ret;
switch (dType) {
case HALF:
ret = Arm_DataType::F16;
break;
case FLOAT32:
ret = Arm_DataType::F32;
break;
case DOUBLE:
ret = Arm_DataType::F64;
break;
case INT8:
ret = Arm_DataType::S8;
break;
case INT16:
ret = Arm_DataType::S16;
break;
case INT32:
ret = Arm_DataType::S32;
break;
case INT64:
ret = Arm_DataType::S64;
break;
case UINT8:
ret = Arm_DataType::U8;
break;
case UINT16:
ret = Arm_DataType::U16;
break;
case UINT32:
ret = Arm_DataType::U32;
break;
case UINT64:
ret = Arm_DataType::U64;
break;
case BFLOAT16:
ret = Arm_DataType::BFLOAT16;
break;
default:
ret = Arm_DataType::UNKNOWN;
};
return ret;
}
bool isArmcomputeFriendly(NDArray& arr) {
auto dType = getArmType(arr.dataType());
int rank = (int)(arr.rankOf());
int ind = arr.ordering() == 'c' ? rank - 1 : 0;
auto arrStrides = arr.stridesOf();
return dType != Arm_DataType::UNKNOWN && rank <= arm_compute::MAX_DIMS && arr.ordering() == 'c' &&
arrStrides[ind] == 1;
}
Arm_TensorInfo getArmTensorInfo(int rank, sd::LongType* bases, sd::DataType ndArrayType,
arm_compute::DataLayout layout) {
constexpr int numChannels = 1;
auto dType = getArmType(ndArrayType);
Arm_TensorShape shape;
shape.set_num_dimensions(rank);
for (int i = 0, j = rank - 1; i < rank; i++, j--) {
shape[i] = static_cast<uint32_t>(bases[j]);
}
// fill the rest unused with 1
for (int i = rank; i < arm_compute::MAX_DIMS; i++) {
shape[i] = 1;
}
return Arm_TensorInfo(shape, numChannels, dType, layout);
}
Arm_TensorInfo getArmTensorInfo(NDArray& arr, arm_compute::DataLayout layout) {
auto dType = getArmType(arr.dataType());
internal_print_nd_shape(arr, "shape");
internal_print_nd_array(arr, "data");
//
constexpr int numChannels = 1;
int rank = (int)(arr.rankOf());
auto bases = arr.shapeOf();
auto arrStrides = arr.stridesOf();
// https://arm-software.github.io/ComputeLibrary/v20.05/_dimensions_8h_source.xhtml
// note: underhood it is stored as std::array<T, num_SD_MAX_DIMENSIONs> _id;
// TensorShape is derived from Dimensions<uint32_t>
// as well as Strides : public Dimensions<uint32_t>
Arm_TensorShape shape;
Arm_Strides strides;
shape.set_num_dimensions(rank);
strides.set_num_dimensions(rank);
size_t element_size = arr.sizeOfT();
for (int i = 0, j = rank - 1; i < rank; i++, j--) {
shape[i] = static_cast<uint32_t>(bases[j]);
strides[i] = static_cast<uint32_t>(arrStrides[j] * element_size);
}
// fill the rest unused with 1
for (int i = rank; i < arm_compute::MAX_DIMS; i++) {
shape[i] = 1;
}
size_t total_size = arr.lengthOf() * element_size;
size_t offset = 0;
if (arr.hasPaddedBuffer()) {
internal_printf("---has padded buffer %d\n", 0);
total_size = arr.getDataBuffer()->getLenInBytes();
offset = arr.offset() * element_size;
}
internal_printf(":: offset %d el size %d arr.getDataBuffer()->getLenInBytes() %d lengthof %d \n",
(int)arr.offset(), (int)element_size, (int)arr.getDataBuffer()->getLenInBytes(),
(int)arr.lengthOf());
Arm_TensorInfo info;
info.init(shape, numChannels, dType, strides, offset, total_size);
info.set_data_layout(layout);
return info;
}
Arm_Tensor getArmTensor(NDArray& arr, arm_compute::DataLayout layout) {
// - Ownership of the backing memory is not transferred to the tensor itself.
// - The tensor mustn't be memory managed.
// - Padding requirements should be accounted by the client code.
// In other words, if padding is required by the tensor after the function
// configuration step, then the imported backing memory should account for it.
// Padding can be checked through the TensorInfo::padding() interface.
// Import existing pointer as backing memory
auto info = getArmTensorInfo(arr, layout);
Arm_Tensor tensor;
tensor.allocator()->init(info);
// get without offset
void* buff = arr.getDataBuffer()->primary();
tensor.allocator()->import_memory(buff);
return tensor;
}
void copyFromTensor(const Arm_Tensor& inTensor, sd::NDArray& output) {
// only for C order
if (output.ordering() != 'c') return;
sd::LongType* shapeInfo = output.shapeInfo();
sd::LongType* bases = &(shapeInfo[1]);
sd::LongType rank = shapeInfo[0];
sd::LongType* strides = output.stridesOf();
int width = bases[rank - 1];
uint8_t* outputBuffer = (uint8_t*)output.buffer();
size_t offset = 0;
arm_compute::Window window;
arm_compute::Iterator tensor_it(&inTensor, window);
int element_size = inTensor.info()->element_size();
window.use_tensor_dimensions(inTensor.info()->tensor_shape(), /* first_dimension =*/arm_compute::Window::DimY);
sd::LongType coords[SD_MAX_RANK] = {};
auto copySize = width * element_size;
arm_compute::execute_window_loop(
window,
[&](const arm_compute::Coordinates& id) {
auto src = tensor_it.ptr();
auto dest = outputBuffer + offset * element_size;
memcpy(dest, src, copySize);
offset = sd::inc_coords(bases, strides, coords, offset, rank, 1);
},
tensor_it);
}
void copyToTensor(sd::NDArray& input, Arm_Tensor& outTensor) {
// only for C order
if (input.ordering() != 'c') return;
sd::LongType* shapeInfo = input.shapeInfo();
sd::LongType* bases = &(shapeInfo[1]);
sd::LongType rank = shapeInfo[0];
sd::LongType* strides = input.stridesOf();
uint8_t* inputBuffer = (uint8_t*)input.buffer();
int width = bases[rank - 1];
size_t offset = 0;
arm_compute::Window window;
arm_compute::Iterator tensor_it(&outTensor, window);
int element_size = outTensor.info()->element_size();
window.use_tensor_dimensions(outTensor.info()->tensor_shape(), /* first_dimension =*/arm_compute::Window::DimY);
sd::LongType coords[SD_MAX_RANK] = {};
auto copySize = width * element_size;
arm_compute::execute_window_loop(
window,
[&](const arm_compute::Coordinates& id) {
auto dest = tensor_it.ptr();
auto src = inputBuffer + offset * element_size;
memcpy(dest, src, copySize);
offset = sd::inc_coords(bases, strides, coords, offset, rank, 1);
},
tensor_it);
}
// armcompute should be built with debug option
void print_tensor(Arm_ITensor& tensor, const char* msg) {
auto info = tensor.info();
auto padding = info->padding();
std::cout << msg << "\ntotal: " << info->total_size() << "\n";
for (int i = 0; i < arm_compute::MAX_DIMS; i++) {
std::cout << info->dimension(i) << ",";
}
std::cout << std::endl;
for (int i = 0; i < arm_compute::MAX_DIMS; i++) {
std::cout << info->strides_in_bytes()[i] << ",";
}
std::cout << "\npadding: l " << padding.left << ", r " << padding.right << ", t " << padding.top << ", b "
<< padding.bottom << std::endl;
#ifdef ARM_COMPUTE_ASSERTS_ENABLED
// note it did not print correctly fro NHWC
std::cout << msg << ":\n";
tensor.print(std::cout);
std::cout << std::endl;
#endif
}
} // namespace platforms
} // namespace ops
} // namespace sd
libnd4j/include/ops/declarable/platform/armcompute/armcomputeUtils.h
+303
View File
@@ -0,0 +1,303 @@
/*
* ******************************************************************************
* *
* *
* * 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
* *****************************************************************************
*/
// Created by Abdelrauf 2020
#ifndef DEV_TESTSARMCOMPUTEUTILS_H
#define DEV_TESTSARMCOMPUTEUTILS_H
#include <arm_compute/core/Helpers.h>
#include <arm_compute/core/ITensor.h>
#include <arm_compute/core/Strides.h>
#include <arm_compute/core/TensorInfo.h>
#include <arm_compute/core/TensorShape.h>
#include <arm_compute/core/Types.h>
#include <arm_compute/core/Validate.h>
#include <arm_compute/core/Window.h>
#include <arm_compute/runtime/NEON/NEFunctions.h>
#include <arm_compute/runtime/Tensor.h>
#include <arm_compute/runtime/TensorAllocator.h>
#include <array/NDArray.h>
#include <graph/Context.h>
#include <legacy/NativeOps.h>
#include <ops/declarable/PlatformHelper.h>
#include <system/platform_boilerplate.h>
#include <iostream>
using namespace samediff;
#if 0
#define internal_printf(FORMAT, ...) sd_printf(FORMAT, __VA_ARGS__)
#define internal_print_arm_array(a, b) print_tensor(a, b)
#define internal_print_nd_array(a, b) ((a).printIndexedBuffer(b))
#define internal_print_nd_shape(a, b) ((a).printShapeInfo(b))
#else
#define internal_printf(FORMAT, ...)
#define internal_print_arm_array(a, b)
#define internal_print_nd_array(a, b)
#define internal_print_nd_shape(a, b)
#endif
namespace sd {
namespace ops {
namespace platforms {
using Arm_DataType = arm_compute::DataType;
using Arm_Tensor = arm_compute::Tensor;
using Arm_ITensor = arm_compute::ITensor;
using Arm_TensorInfo = arm_compute::TensorInfo;
using Arm_TensorShape = arm_compute::TensorShape;
using Arm_Strides = arm_compute::Strides;
using Arm_WeightsInfo = arm_compute::WeightsInfo;
using Arm_PermutationVector = arm_compute::PermutationVector;
using Arm_DataLayout = arm_compute::DataLayout;
/**
* Here we actually declare our platform helpers
*/
DECLARE_PLATFORM(maxpool2d, ENGINE_CPU);
DECLARE_PLATFORM(avgpool2d, ENGINE_CPU);
DECLARE_PLATFORM(conv2d, ENGINE_CPU);
DECLARE_PLATFORM(deconv2d, ENGINE_CPU);
// utils
Arm_DataType getArmType(const sd::DataType& dType);
Arm_TensorInfo getArmTensorInfo(int rank, sd::LongType* bases, sd::DataType ndArrayType,
Arm_DataLayout layout = Arm_DataLayout::UNKNOWN);
Arm_TensorInfo getArmTensorInfo(NDArray& arr, Arm_DataLayout layout = Arm_DataLayout::UNKNOWN);
Arm_Tensor getArmTensor(NDArray& arr, Arm_DataLayout layout = Arm_DataLayout::UNKNOWN);
void copyFromTensor(const Arm_Tensor& inTensor, NDArray& output);
void copyToTensor(NDArray& input, Arm_Tensor& outTensor);
void print_tensor(Arm_ITensor& tensor, const char* msg);
bool isArmcomputeFriendly(NDArray& arr);
template <typename F>
class ArmFunction {
public:
template <typename... Args>
void configure(NDArray* input, NDArray* output, Arm_DataLayout layout, Args&&... args) {
bool inputHasPaddedBuffer = input->hasPaddedBuffer();
bool outputHasPaddedBuffer = output->hasPaddedBuffer();
if (inputHasPaddedBuffer) {
in = getArmTensor(*input, layout);
internal_printf("input is a padded buffer %d\n", 0);
} else {
auto inInfo = getArmTensorInfo(*input, layout);
in.allocator()->init(inInfo);
}
if (outputHasPaddedBuffer) {
out = getArmTensor(*output, layout);
internal_printf("output is a padded buffer %d\n", 0);
} else {
auto outInfo = getArmTensorInfo(*output, layout);
out.allocator()->init(outInfo);
}
armFunction.configure(&in, &out, std::forward<Args>(args)...);
if (!inputHasPaddedBuffer) {
if (in.info()->has_padding()) {
// allocate and copy
in.allocator()->allocate();
inputNd = input;
} else {
// import only for ews()==1
in.allocator()->import_memory(input->buffer());
internal_printf("input import %d\n", 0);
}
}
if (!outputHasPaddedBuffer) {
if (out.info()->has_padding()) {
// store pointer to our array to copy after run
out.allocator()->allocate();
outNd = output;
} else {
// import only for ews()==1
out.allocator()->import_memory(output->buffer());
}
}
}
void run() {
if (inputNd) {
// copy
copyToTensor(*inputNd, in);
}
armFunction.run();
if (outNd) {
copyFromTensor(out, *outNd);
internal_printf("output copy %d\n", 0);
internal_print_arm_array(out, "out");
}
}
private:
Arm_Tensor in;
Arm_Tensor out;
NDArray* inputNd = nullptr;
NDArray* outNd = nullptr;
F armFunction{};
};
template <typename F>
class ArmFunctionWeighted {
public:
template <typename... Args>
void configure(NDArray* input, NDArray* weights, NDArray* biases, NDArray* output, Arm_DataLayout layout,
arm_compute::PermutationVector permuteVector, Args&&... args) {
bool inputHasPaddedBuffer = input->hasPaddedBuffer();
bool weightsHasPaddedBuffer = weights->hasPaddedBuffer();
bool outputHasPaddedBuffer = output->hasPaddedBuffer();
bool biasesHasPaddedBuffer = false;
if (inputHasPaddedBuffer) {
in = getArmTensor(*input, layout);
} else {
in.allocator()->init(getArmTensorInfo(*input, layout));
}
if (weightsHasPaddedBuffer) {
w = getArmTensor(*weights, layout);
} else {
w.allocator()->init(getArmTensorInfo(*weights, layout));
}
if (outputHasPaddedBuffer) {
out = getArmTensor(*output, layout);
} else {
out.allocator()->init(getArmTensorInfo(*output, layout));
}
Arm_Tensor* bias_ptr = nullptr;
if (biases) {
biasesHasPaddedBuffer = biases->hasPaddedBuffer();
if (biasesHasPaddedBuffer) {
b = getArmTensor(*biases, layout);
} else {
b.allocator()->init(getArmTensorInfo(*biases, layout));
}
bias_ptr = &b;
}
if (permuteVector.num_dimensions() == 0) {
armFunction.configure(&in, &w, bias_ptr, &out, std::forward<Args>(args)...);
} else {
// configure with permute kernel
Arm_TensorShape shape;
int rank = permuteVector.num_dimensions();
shape.set_num_dimensions(rank);
auto wInfoPtr = w.info();
for (int i = 0; i < rank; i++) {
shape[i] = wInfoPtr->dimension(permuteVector[i]);
}
for (int i = rank; i < arm_compute::MAX_DIMS; i++) {
shape[i] = 1;
}
Arm_TensorInfo wPermInfo(shape, 1, wInfoPtr->data_type(), layout);
wPerm.allocator()->init(wPermInfo);
permuter.configure(&w, &wPerm, permuteVector);
armFunction.configure(&in, &wPerm, bias_ptr, &out, std::forward<Args>(args)...);
wPerm.allocator()->allocate();
runPerm = true;
}
// import buffer
if (!inputHasPaddedBuffer) {
if (in.info()->has_padding()) {
// allocate and copy
in.allocator()->allocate();
inputNd = input;
} else {
// import buffer
in.allocator()->import_memory(input->buffer());
}
}
if (!weightsHasPaddedBuffer) {
if (w.info()->has_padding()) {
// store pointer to our array to copy after run
w.allocator()->allocate();
wNd = weights;
} else {
// import
w.allocator()->import_memory(weights->buffer());
}
}
if (biases && !biasesHasPaddedBuffer) {
if (b.info()->has_padding()) {
// store pointer to our array to copy after run
b.allocator()->allocate();
bNd = biases;
} else {
// import
b.allocator()->import_memory(biases->buffer());
}
}
if (!outputHasPaddedBuffer) {
if (out.info()->has_padding()) {
// store pointer to our array to copy after run
out.allocator()->allocate();
outNd = output;
} else {
// import
out.allocator()->import_memory(output->buffer());
}
}
}
void run() {
if (inputNd) {
// copy
copyToTensor(*inputNd, in);
}
if (bNd) {
// copy
copyToTensor(*bNd, b);
}
if (wNd) {
// copy
copyToTensor(*wNd, w);
}
if (runPerm) {
permuter.run();
}
armFunction.run();
if (outNd) {
copyFromTensor(out, *outNd);
}
}
private:
bool runPerm = false;
Arm_Tensor in;
Arm_Tensor b;
Arm_Tensor w;
Arm_Tensor wPerm;
Arm_Tensor out;
NDArray* inputNd = nullptr;
NDArray* wNd = nullptr;
NDArray* bNd = nullptr;
NDArray* outNd = nullptr;
arm_compute::NEPermute permuter;
F armFunction{};
};
} // namespace platforms
} // namespace ops
} // namespace sd
#endif // DEV_TESTSARMCOMPUTEUTILS_H
libnd4j/include/ops/declarable/platform/armcompute/avgpooling2d.cpp
+115
View File
@@ -0,0 +1,115 @@
/*
* ******************************************************************************
* *
* *
* * 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
* *****************************************************************************
*/
// Created by Abdelrauf (rauf@konduit.ai) 2020
#include <ops/declarable/OpRegistrator.h>
#include <ops/declarable/PlatformHelper.h>
#include <ops/declarable/helpers/convolutions.h>
#include <system/platform_boilerplate.h>
#include "armcomputeUtils.h"
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 ARMCOMPUTE op: input should have rank of 4, but got %i instead",
input->rankOf());
REQUIRE_TRUE(dH != 0 && dW != 0, 0, "AVGPOOL2D ARMCOMPUTE op: dilation must not be zero, but got instead {%i, %i}",
dH, dW);
bool excludePadding = (extraParam0 == 0) ? true : false;
auto dataLayout = isNCHW ? arm_compute::DataLayout::NCHW : arm_compute::DataLayout::NHWC;
// Calculate individual paddings
sd::LongType padLeft, padTop, padRight, padBottom;
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);
}
padLeft = pW;
padTop = pH;
padRight = (oW - 1) * sW - iW + kW - pW;
padBottom = (oH - 1) * sH - iH + kH - pH;
auto poolPad = arm_compute::PadStrideInfo(sW, sH, padLeft, padRight, padTop, padBottom,
arm_compute::DimensionRoundingType::FLOOR);
auto poolInfo = arm_compute::PoolingLayerInfo(arm_compute::PoolingType::AVG, arm_compute::Size2D(kW, kH), dataLayout,
poolPad, excludePadding);
ArmFunction<arm_compute::NEPoolingLayer> pool;
pool.configure(input, output, dataLayout, poolInfo);
pool.run(); // run function
return sd::Status::OK;
}
//////////////////////////////////////////////////////////////////////////
PLATFORM_CHECK(avgpool2d, ENGINE_CPU) {
auto input = INPUT_VARIABLE(0);
auto output = OUTPUT_VARIABLE(0);
const sd::LongType dH = INT_ARG(6);
const sd::LongType dW = INT_ARG(7);
// Data types supported: QASYMM8/QASYMM8_SIGNED/F16/F32
// for now, we will ignore F16 as it shoulde be preconditioned for pool size 2,3 and arm64-v8.2-a architecture
Requirements req("ARMCOMPUTE AVGPOOL2d OP");
req.expectEq(makeInfoVariable(input->dataType(), TYPE_MSG_INPUT), DataType::FLOAT32) &&
req.expectEq(makeInfoVariable(output->dataType(), TYPE_MSG_OUTPUT), DataType::FLOAT32) &&
req.expectEq(makeInfoVariable(dH, "dilation#H"), 1) && req.expectEq(makeInfoVariable(dW, "dilation#W"), 1) &&
req.expectLessEq(makeInfoVariable(input->rankOf(), RANK_MSG_INPUT), arm_compute::MAX_DIMS) &&
req.expectEq(makeInfoVariable(input->ordering(), ORDERING_MSG_INPUT), 'c') &&
req.expectEq(makeInfoVariable(input->stridesOf()[input->rankOf() - 1], "input#lastStride"), 1) &&
req.expectLessEq(makeInfoVariable(output->rankOf(), RANK_MSG_OUTPUT), arm_compute::MAX_DIMS) &&
req.expectEq(makeInfoVariable(output->ordering(), ORDERING_MSG_OUTPUT), 'c') &&
req.expectEq(makeInfoVariable(output->stridesOf()[output->rankOf() - 1], "output#lastStride"), 1);
req.logTheSuccess();
return req;
}
} // namespace platforms
} // namespace ops
} // namespace sd
libnd4j/include/ops/declarable/platform/armcompute/conv2d.cpp
+148
View File
@@ -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
* *****************************************************************************
*/
// Created by Abdelrauf (rauf@konduit.ai) 2020
#include <ops/declarable/OpRegistrator.h>
#include <ops/declarable/PlatformHelper.h>
#include <ops/declarable/helpers/convolutions.h>
#include <system/platform_boilerplate.h>
#include "armcomputeUtils.h"
namespace sd {
namespace ops {
namespace platforms {
//////////////////////////////////////////////////////////////////////
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
// Calculate individual paddings
sd::LongType padLeft, padTop, padRight, padBottom;
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);
ConvolutionUtils::calcPadding2D(pH, pW, oH, oW, iH, iW, kH, kW, sH, sW, dH, dW, paddingMode);
int pWSame = (paddingMode == 2 && dW > 1) ? ((oW - 1) * sW + (kW - 1) * dW + 1 - iW) / 2
: pW; // dH == 1 for causal mode in conv1d
padLeft = pW;
padTop = pH;
padRight = (oW - 1) * sW - iW + kW - pWSame;
padBottom = (oH - 1) * sH - iH + kH - pH;
std::vector<sd::LongType> expectedWeightsShape = ConvolutionUtils::expectWeightsShape(wFormat, kH, kW, iC, oC);
REQUIRE_TRUE(weights->isSameShape(expectedWeightsShape), 0,
"CONV2D ARMCOMPUTE 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 ARMCOMPUTE OP: wrong shape of array with biases, expected rank, length: <=2, %i, but got %i, "
"%i instead !",
oC, bias->rankOf(), bias->lengthOf());
auto dataLayout = isNCHW ? arm_compute::DataLayout::NCHW : arm_compute::DataLayout::NHWC;
// check weight input datalayout match
bool dataLayoutMatch = (isNCHW && wFormat == 1) || (!isNCHW && wFormat == 2);
arm_compute::PermutationVector permuteVector;
if (!dataLayoutMatch) {
// lets premute
if (wFormat == 0) {
if (isNCHW) {
// reshape
permuteVector = arm_compute::PermutationVector(2U, 3U, 1U, 0U);
} else {
// reshape
permuteVector = arm_compute::PermutationVector(1U, 2U, 3U, 0U);
}
} else if (wFormat == 1) {
permuteVector = arm_compute::PermutationVector(2U, 0U, 1U, 3U);
} else {
permuteVector = arm_compute::PermutationVector(1U, 2U, 0U, 3U);
}
} else {
// set 0
permuteVector.set_num_dimensions(0);
}
Arm_WeightsInfo wInfo(false, kW, kH, 1);
arm_compute::Size2D dilation(dW, dH);
arm_compute::PadStrideInfo pad(sW, sH, padLeft, padRight, padTop, padBottom,
arm_compute::DimensionRoundingType::FLOOR);
ArmFunctionWeighted<arm_compute::NEConvolutionLayer> conv;
conv.configure(input, weights, bias, output, dataLayout, permuteVector, pad, wInfo, dilation);
conv.run(); // run function
return sd::Status::OK;
}
PLATFORM_CHECK(conv2d, ENGINE_CPU) {
auto input = INPUT_VARIABLE(0);
auto weights = INPUT_VARIABLE(1);
auto output = OUTPUT_VARIABLE(0);
// Data types supported: QASYMM8/QASYMM8_SIGNED/F16/F32.
Requirements req("ARMCOMPUTE CONV2d OP");
req.expectEq(makeInfoVariable(input->dataType(), TYPE_MSG_INPUT0), DataType::FLOAT32) &&
req.expectEq(makeInfoVariable(weights->dataType(), TYPE_MSG_INPUT1), DataType::FLOAT32) &&
req.expectEq(makeInfoVariable(output->dataType(), TYPE_MSG_OUTPUT), DataType::FLOAT32) &&
req.expectLessEq(makeInfoVariable(input->rankOf(), RANK_MSG_INPUT0), arm_compute::MAX_DIMS) &&
req.expectEq(makeInfoVariable(input->ordering(), ORDERING_MSG_INPUT0), 'c') &&
req.expectEq(makeInfoVariable(input->stridesOf()[input->rankOf() - 1], "input0#lastStride"), 1) &&
req.expectLessEq(makeInfoVariable(weights->rankOf(), RANK_MSG_INPUT1), arm_compute::MAX_DIMS) &&
req.expectEq(makeInfoVariable(weights->ordering(), ORDERING_MSG_INPUT1), 'c') &&
req.expectEq(makeInfoVariable(weights->stridesOf()[weights->rankOf() - 1], "input1#lastStride"), 1) &&
req.expectLessEq(makeInfoVariable(output->rankOf(), RANK_MSG_OUTPUT), arm_compute::MAX_DIMS) &&
req.expectEq(makeInfoVariable(output->ordering(), ORDERING_MSG_OUTPUT), 'c') &&
req.expectEq(makeInfoVariable(output->stridesOf()[output->rankOf() - 1], "output#lastStride"), 1);
req.logTheSuccess();
return req;
}
} // namespace platforms
} // namespace ops
} // namespace sd
libnd4j/include/ops/declarable/platform/armcompute/deconv2d.cpp
+159
View File
@@ -0,0 +1,159 @@
/*
* ******************************************************************************
* *
* *
* * 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
* *****************************************************************************
*/
// Created by Abdelrauf (rauf@konduit.ai) 2020
#include <ops/declarable/OpRegistrator.h>
#include <ops/declarable/PlatformHelper.h>
#include <ops/declarable/helpers/convolutions.h>
#include <system/platform_boilerplate.h>
#include "armcomputeUtils.h"
namespace sd {
namespace ops {
namespace platforms {
//////////////////////////////////////////////////////////////////////
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 ARMCOMPUTE OP: rank of input array must be equal to 4, but got %i instead !",
input->rankOf());
REQUIRE_TRUE(weights->rankOf() == 4, 0,
"CUSTOM DECONV2D ARMCOMPUTE 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
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]
// Calculate individual paddings
sd::LongType padLeft, padTop, padRight, padBottom;
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 ARMCOMPUTE 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 ARMCOMPUTE OP: wrong shape of array with biases, expected rank, length: <=2, %i, but "
"got %i, %i instead !",
oC, bias->rankOf(), bias->lengthOf());
if (paddingMode) {
// 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);
}
padLeft = pW;
padTop = pH;
padRight = (iW - 1) * sW - oW + kW - pW;
padBottom = (iH - 1) * sH - oH + kH - pH;
auto dataLayout = isNCHW ? arm_compute::DataLayout::NCHW : arm_compute::DataLayout::NHWC;
// check weight input datalayout match
bool dataLayoutMatch = (isNCHW && wFormat == 1) || (!isNCHW && wFormat == 2);
arm_compute::PermutationVector permuteVector;
// unlike in cov2d for weights iC and oC permutted : for example {oC, iC, kH, kW}, {iC, oC, kH, kW}
// but we need it normal way for arm
if (!dataLayoutMatch) {
// lets premute
if (wFormat == 0) {
if (isNCHW) {
// reshape
permuteVector = arm_compute::PermutationVector(2U, 3U, 0U, 1U);
} else {
// reshape
permuteVector = arm_compute::PermutationVector(0U, 2U, 3U, 1U);
}
} else if (wFormat == 1) {
permuteVector = arm_compute::PermutationVector(3U, 0U, 1U, 2U);
} else {
permuteVector = arm_compute::PermutationVector(1U, 2U, 3U, 0U);
}
} else {
// fix weight
if (isNCHW) {
permuteVector = arm_compute::PermutationVector(0U, 1U, 3U, 2U);
} else {
permuteVector = arm_compute::PermutationVector(3U, 1U, 2U, 0U);
}
}
Arm_WeightsInfo wInfo(false, kW, kH, 1);
arm_compute::PadStrideInfo pad(sW, sH, padLeft, padRight, padTop, padBottom,
arm_compute::DimensionRoundingType::FLOOR);
ArmFunctionWeighted<arm_compute::NEDeconvolutionLayer> deconv;
deconv.configure(input, weights, bias, output, dataLayout, permuteVector, pad);
deconv.run(); // run function
return sd::Status::OK;
}
PLATFORM_CHECK(deconv2d, ENGINE_CPU) {
auto input = INPUT_VARIABLE(0);
auto weights = INPUT_VARIABLE(1);
auto output = OUTPUT_VARIABLE(0);
int dH = INT_ARG(6);
int dW = INT_ARG(7);
// Data types supported: QASYMM8/QASYMM8_SIGNED/F16/F32.
Requirements req("ARMCOMPUTE DECONV2d OP");
req.expectEq(makeInfoVariable(input->dataType(), TYPE_MSG_INPUT0), DataType::FLOAT32) &&
req.expectEq(makeInfoVariable(weights->dataType(), TYPE_MSG_INPUT1), DataType::FLOAT32) &&
req.expectEq(makeInfoVariable(output->dataType(), TYPE_MSG_OUTPUT), DataType::FLOAT32) &&
req.expectEq(makeInfoVariable(dH, "dilation#H"), 1) && req.expectEq(makeInfoVariable(dW, "dilation#W"), 1) &&
req.expectLessEq(makeInfoVariable(input->rankOf(), RANK_MSG_INPUT0), arm_compute::MAX_DIMS) &&
req.expectEq(makeInfoVariable(input->ordering(), ORDERING_MSG_INPUT0), 'c') &&
req.expectEq(makeInfoVariable(input->stridesOf()[input->rankOf() - 1], "input0#lastStride"), 1) &&
req.expectLessEq(makeInfoVariable(weights->rankOf(), RANK_MSG_INPUT1), arm_compute::MAX_DIMS) &&
req.expectEq(makeInfoVariable(weights->ordering(), ORDERING_MSG_INPUT1), 'c') &&
req.expectEq(makeInfoVariable(weights->stridesOf()[weights->rankOf() - 1], "input1#lastStride"), 1) &&
req.expectLessEq(makeInfoVariable(output->rankOf(), RANK_MSG_OUTPUT), arm_compute::MAX_DIMS) &&
req.expectEq(makeInfoVariable(output->ordering(), ORDERING_MSG_OUTPUT), 'c') &&
req.expectEq(makeInfoVariable(output->stridesOf()[output->rankOf() - 1], "output#lastStride"), 1);
req.logTheSuccess();
return req;
}
} // namespace platforms
} // namespace ops
} // namespace sd
libnd4j/include/ops/declarable/platform/armcompute/maxpooling2d.cpp
+111
View File
@@ -0,0 +1,111 @@
/*
* ******************************************************************************
* *
* *
* * 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
* *****************************************************************************
*/
// Created by Abdelrauf 2020
#include <ops/declarable/OpRegistrator.h>
#include <ops/declarable/PlatformHelper.h>
#include <ops/declarable/helpers/convolutions.h>
#include <system/platform_boilerplate.h>
#include "armcomputeUtils.h"
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 ARMCOMPUTE 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
auto dataLayout = isNCHW ? arm_compute::DataLayout::NCHW : arm_compute::DataLayout::NHWC;
// Calculate individual paddings
sd::LongType padLeft, padTop, padRight, padBottom;
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);
}
padLeft = pW;
padTop = pH;
padRight = (oW - 1) * sW - iW + kW - pW;
padBottom = (oH - 1) * sH - iH + kH - pH;
auto poolPad = arm_compute::PadStrideInfo(sW, sH, padLeft, padRight, padTop, padBottom,
arm_compute::DimensionRoundingType::FLOOR);
auto poolInfo =
arm_compute::PoolingLayerInfo(arm_compute::PoolingType::MAX, arm_compute::Size2D(kW, kH), dataLayout, poolPad);
ArmFunction<arm_compute::NEPoolingLayer> pool;
pool.configure(input, output, dataLayout, poolInfo);
pool.run(); // run function
return sd::Status::OK;
}
//////////////////////////////////////////////////////////////////////////
PLATFORM_CHECK(maxpool2d, ENGINE_CPU) {
auto input = INPUT_VARIABLE(0);
auto output = OUTPUT_VARIABLE(0);
const int dH = INT_ARG(6);
const int dW = INT_ARG(7);
// Data types supported: QASYMM8/QASYMM8_SIGNED/F16/F32
// for now, we will ignore F16 as it shoulde be preconditioned for pool size 2,3 and arm64-v8.2-a architecture
Requirements req("ARMCOMPUTE MAXPOOL2d OP");
req.expectEq(makeInfoVariable(input->dataType(), TYPE_MSG_INPUT), DataType::FLOAT32) &&
req.expectEq(makeInfoVariable(output->dataType(), TYPE_MSG_OUTPUT), DataType::FLOAT32) &&
req.expectEq(makeInfoVariable(dH, "dilation#H"), 1) && req.expectEq(makeInfoVariable(dW, "dilation#W"), 1) &&
req.expectLessEq(makeInfoVariable(input->rankOf(), RANK_MSG_INPUT), arm_compute::MAX_DIMS) &&
req.expectEq(makeInfoVariable(input->ordering(), ORDERING_MSG_INPUT), 'c') &&
req.expectEq(makeInfoVariable(input->stridesOf()[input->rankOf() - 1], "input#lastStride"), 1) &&
req.expectLessEq(makeInfoVariable(output->rankOf(), RANK_MSG_OUTPUT), arm_compute::MAX_DIMS) &&
req.expectEq(makeInfoVariable(output->ordering(), ORDERING_MSG_OUTPUT), 'c') &&
req.expectEq(makeInfoVariable(output->stridesOf()[output->rankOf() - 1], "output#lastStride"), 1);
req.logTheSuccess();
return req;
}
} // namespace platforms
} // namespace ops
} // namespace sd
libnd4j/include/ops/declarable/platform/cudnn/avgpool2d.cu
+164
View File
@@ -0,0 +1,164 @@
/* ******************************************************************************
*
*
* This program and the accompanying materials are made available under the
* terms of the Apache License, Version 2.0 which is available at
* https://www.apache.org/licenses/LICENSE-2.0.
*
* See the NOTICE file distributed with this work for additional
* information regarding copyright ownership.
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
* WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
* License for the specific language governing permissions and limitations
* under the License.
*
* SPDX-License-Identifier: Apache-2.0
******************************************************************************/
//
// @author Yurii Shyrma (iuriish@yahoo.com)
//
#include <ops/declarable/helpers/convolutions.h>
#include "cudnnUtils.h"
namespace sd {
namespace ops {
namespace platforms {
//////////////////////////////////////////////////////////////////////////
PLATFORM_IMPL(avgpool2d, ENGINE_CUDA) {
auto input = INPUT_VARIABLE(0);
auto output = OUTPUT_VARIABLE(0);
// 0,1 - kernel Height/Width; 2,3 - stride Height/Width; 4,5 - pad Height/Width; 6,7 - dilation Height/Width; 8 - same
// mode;
const LongType kH = INT_ARG(0);
const LongType kW = INT_ARG(1);
const LongType sH = INT_ARG(2);
const LongType sW = INT_ARG(3);
LongType pH = INT_ARG(4);
LongType pW = INT_ARG(5);
const LongType dH = INT_ARG(6);
const LongType dW = INT_ARG(7);
const auto paddingMode = static_cast<bool>(INT_ARG(8));
const auto extraParam0 = INT_ARG(9);
const int isNCHW = block.getIArguments()->size() > 10 ? !INT_ARG(10) : 1; // INT_ARG(10): 0-NCHW, 1-NHWC
REQUIRE_TRUE(input->rankOf() == 4, 0, "AVGPOOL2D CUDNN op: input should have rank of 4, but got %i instead",
input->rankOf());
REQUIRE_TRUE(dH != 0 && dW != 0, 0, "AVGPOOL2D CUDNN op: dilation must not be zero, but got instead {%i, %i}", dH,
dW);
LongType oH = 0;
LongType oW = 0;
const LongType iH = static_cast<LongType>(isNCHW ? input->sizeAt(2) : input->sizeAt(1));
const LongType iW = static_cast<LongType>(isNCHW ? input->sizeAt(3) : input->sizeAt(2));
ConvolutionUtils::calcOutSizePool2D(oH, oW, kH, kW, sH, sW, pH, pW, dH, dW, iH, iW, paddingMode);
if (paddingMode) ConvolutionUtils::calcPadding2D(pH, pW, oH, oW, iH, iW, kH, kW, sH, sW, dH, dW);
const cudnnPoolingMode_t mode =
(extraParam0 == 0) ? CUDNN_POOLING_AVERAGE_COUNT_EXCLUDE_PADDING : CUDNN_POOLING_AVERAGE_COUNT_INCLUDE_PADDING;
pooling2dCUDNN(block.launchContext(), input, output, kH, kW, sH, sW, pH, pW, dH, dW, isNCHW, mode);
return Status::OK;
}
//////////////////////////////////////////////////////////////////////////
PLATFORM_CHECK(avgpool2d, ENGINE_CUDA) {
auto input = INPUT_VARIABLE(0);
auto output = OUTPUT_VARIABLE(0);
Requirements req("CUDNN AVGPOOL2d OP");
req.expectEq(makeInfoVariable(input->dataType(), TYPE_MSG_INPUT),
makeInfoVariable(output->dataType(), TYPE_MSG_OUTPUT)) &&
req.expectIn(makeInfoVariable(input->dataType(), TYPE_MSG_INPUT),
{INT32, HALF, FLOAT32, DOUBLE});
req.logTheSuccess();
return req;
}
//////////////////////////////////////////////////////////////////////////
PLATFORM_IMPL(avgpool2d_bp, ENGINE_CUDA) {
auto input = INPUT_VARIABLE(0); // [bS, iH, iW, iC] (NHWC) or [bS, iC, iH, iW] (NCHW)
auto gradO = INPUT_VARIABLE(1); // [bS, oH, oW, oC] (NHWC) or [bS, oC, oH, oW] (NCHW), epsilon_next
auto gradI = OUTPUT_VARIABLE(0); // [bS, iH, iW, iC] (NHWC) or [bS, iC, iH, iW] (NCHW), epsilon
const LongType kH = INT_ARG(0); // filter(kernel) height
const LongType kW = INT_ARG(1); // filter(kernel) width
const LongType sH = INT_ARG(2); // strides height
const LongType sW = INT_ARG(3); // strides width
LongType pH = INT_ARG(4); // paddings height
LongType pW = INT_ARG(5); // paddings width
const LongType dH = INT_ARG(6); // dilations height
const LongType dW = INT_ARG(7); // dilations width
const auto paddingMode = INT_ARG(8); // 0-VALID, 1-SAME
const auto extraParam0 = INT_ARG(9);
const auto isNCHW = block.getIArguments()->size() > 10 ? !INT_ARG(10) : 1; // INT_ARG(10): 0-NCHW, 1-NHWC
REQUIRE_TRUE(input->rankOf() == 4, 0, "AVGPOOL2D_BP CUDNN op: input should have rank of 4, but got %i instead",
input->rankOf());
REQUIRE_TRUE(dH != 0 && dW != 0, 0, "AVGPOOL2D_BP CUDNN op: dilation must not be zero, but got instead {%i, %i}", dH,
dW);
LongType bS, iC, iH, iW, oC, oH,
oW; // batch size, input channels, input height/width, output channels, output height/width;
LongType indIOioC, indIiH, indWoC, indWiC, indWkH, indOoH; // corresponding indexes
ConvolutionUtils::getSizesAndIndexesConv2d(isNCHW, 0, *input, *gradO, bS, iC, iH, iW, oC, oH, oW, indIOioC, indIiH,
indWiC, indWoC, indWkH, indOoH);
std::vector<LongType> expectedGradOShape =
ShapeUtils::composeShapeUsingDimsAndIdx({bS, iC, oH, oW, 0, indIOioC, indIiH, indIiH + 1});
std::vector<LongType> expectedGradIShape =
ShapeUtils::composeShapeUsingDimsAndIdx({bS, iC, iH, iW, 0, indIOioC, indIiH, indIiH + 1});
REQUIRE_TRUE(gradO->isSameShape(expectedGradOShape), 0,
"AVGPOOL2D_BP CUDNN op: wrong shape of output's gradients array (next epsilon), expected is %s, but got "
"%s instead !",
ShapeUtils::shapeAsString(expectedGradOShape).c_str(), ShapeUtils::shapeAsString(gradO).c_str());
REQUIRE_TRUE(
gradI->isSameShape(expectedGradIShape), 0,
"AVGPOOL2D_BP CUDNN op: wrong shape of input's gradients array (epsilon), expected is %s, but got %s instead !",
ShapeUtils::shapeAsString(expectedGradIShape).c_str(), ShapeUtils::shapeAsString(gradI).c_str());
if (paddingMode) // SAME
ConvolutionUtils::calcPadding2D(pH, pW, oH, oW, iH, iW, kH, kW, sH, sW, dH, dW);
const cudnnPoolingMode_t mode =
(extraParam0 == 0) ? CUDNN_POOLING_AVERAGE_COUNT_EXCLUDE_PADDING : CUDNN_POOLING_AVERAGE_COUNT_INCLUDE_PADDING;
pooling2dBpCUDNN(block.launchContext(), input, gradO, gradI, kH, kW, sH, sW, pH, pW, dH, dW, isNCHW, mode);
return Status::OK;
}
PLATFORM_CHECK(avgpool2d_bp, ENGINE_CUDA) {
auto input = INPUT_VARIABLE(0); // [bS, iH, iW, iC] (NHWC) or [bS, iC, iH, iW] (NCHW)
auto gradO = INPUT_VARIABLE(1); // [bS, oH, oW, oC] (NHWC) or [bS, oC, oH, oW] (NCHW), epsilon_next
auto gradI = OUTPUT_VARIABLE(0); // [bS, iH, iW, iC] (NHWC) or [bS, iC, iH, iW] (NCHW), epsilon
Requirements req("CUDNN AVGPOOL2d_BP OP");
req.expectEq(makeInfoVariable(input->dataType(), TYPE_MSG_INPUT0),
makeInfoVariable(gradO->dataType(), TYPE_MSG_INPUT1)) &&
req.expectEq(makeInfoVariable(input->dataType(), TYPE_MSG_INPUT),
makeInfoVariable(gradI->dataType(), TYPE_MSG_OUTPUT)) &&
req.expectIn(makeInfoVariable(input->dataType(), TYPE_MSG_INPUT),
{INT32, HALF, FLOAT32, DOUBLE}) &&
req.expect(
makeShapeInfoVariable(input, SHAPE_MSG_INPUT0), makeShapeInfoVariable(gradI, SHAPE_MSG_OUTPUT),
[](const decltype(input)& l, const decltype(gradI)& r) {
return shape::haveSameShapeAndStrides(l->shapeInfo(), r->shapeInfo());
},
EXPECTED_EQ_MSG);
req.logTheSuccess();
return req;
}
} // namespace platforms
} // namespace ops
} // namespace sd
libnd4j/include/ops/declarable/platform/cudnn/avgpool3d.cu
+176
View File
@@ -0,0 +1,176 @@
/* ******************************************************************************
*
*
* This program and the accompanying materials are made available under the
* terms of the Apache License, Version 2.0 which is available at
* https://www.apache.org/licenses/LICENSE-2.0.
*
* See the NOTICE file distributed with this work for additional
* information regarding copyright ownership.
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
* WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
* License for the specific language governing permissions and limitations
* under the License.
*
* SPDX-License-Identifier: Apache-2.0
******************************************************************************/
//
// @author Yurii Shyrma (iuriish@yahoo.com)
//
#include <ops/declarable/helpers/convolutions.h>
#include "cudnnUtils.h"
namespace sd {
namespace ops {
namespace platforms {
//////////////////////////////////////////////////////////////////////////
PLATFORM_IMPL(avgpool3dnew, ENGINE_CUDA) {
auto input = INPUT_VARIABLE(0); // [bS, iD, iH, iW, iC] (NDHWC) or [bS, iC, iD, iH, iW] (NCDHW)
auto output = OUTPUT_VARIABLE(0); // [bS, oD, oH, oW, iC] (NDHWC) or [bS, iC, oD, oH, oW] (NCDHW)
LongType kD = INT_ARG(0); // filter(kernel) depth
LongType kH = INT_ARG(1); // filter(kernel) height
LongType kW = INT_ARG(2); // filter(kernel) width
LongType sD = INT_ARG(3); // strides depth
LongType sH = INT_ARG(4); // strides height
LongType sW = INT_ARG(5); // strides width
LongType pD = INT_ARG(6); // paddings depth
LongType pH = INT_ARG(7); // paddings height
LongType pW = INT_ARG(8); // paddings width
LongType dD = INT_ARG(9); // dilations depth
LongType dH = INT_ARG(10); // dilations height
LongType dW = INT_ARG(11); // dilations width
int paddingMode = INT_ARG(12); // 1-SAME, 0-VALID
int extraParam0 = INT_ARG(13);
int isNCDHW = block.getIArguments()->size() > 14 ? !INT_ARG(14) : 1; // 0-NCDHW, 1-NDHWC
REQUIRE_TRUE(input->rankOf() == 5, 0,
"AVGPOOL3DNEW CUDNN OP: rank of input array must be equal to 5, but got %i instead !", input->rankOf());
REQUIRE_TRUE(dD != 0 && dH != 0 && dW != 0, 0,
"AVGPOOL3DNEW CUDNN OP: dilation must not be zero, but got instead {%i, %i, %i}", dD, dH, dW);
LongType bS, iC, iD, iH, iW, oC, oD, oH,
oW; // batch size, input channels, input depth/height/width, output channels, output depth/height/width;
LongType indIOioC, indIOioD, indWoC, indWiC, indWkD; // corresponding indexes
ConvolutionUtils::getSizesAndIndexesConv3d(isNCDHW, 0, *input, *output, bS, iC, iD, iH, iW, oC, oD, oH, oW, indIOioC,
indIOioD, indWiC, indWoC, indWkD);
std::vector<LongType> expectedOutputShape =
ShapeUtils::composeShapeUsingDimsAndIdx({bS, iC, oD, oH, oW, 0, indIOioC, indIOioD, indIOioD + 1, indIOioD + 2});
REQUIRE_TRUE(output->isSameShape(expectedOutputShape), 0,
"AVGPOOL3DNEW CUDNN OP: wrong shape of output array, expected is %s, but got %s instead !",
ShapeUtils::shapeAsString(expectedOutputShape).c_str(), ShapeUtils::shapeAsString(output).c_str());
if (paddingMode) // SAME
ConvolutionUtils::calcPadding3D(pD, pH, pW, oD, oH, oW, iD, iH, iW, kD, kH, kW, sD, sH, sW, dD, dH, dW);
const cudnnPoolingMode_t mode =
(extraParam0 == 0) ? CUDNN_POOLING_AVERAGE_COUNT_EXCLUDE_PADDING : CUDNN_POOLING_AVERAGE_COUNT_INCLUDE_PADDING;
pooling3dCUDNN(block.launchContext(), input, output, kD, kH, kW, sD, sH, sW, pD, pH, pW, dD, dH, dW, isNCDHW, mode);
return Status::OK;
}
//////////////////////////////////////////////////////////////////////////
PLATFORM_CHECK(avgpool3dnew, ENGINE_CUDA) {
auto input = INPUT_VARIABLE(0);
auto output = OUTPUT_VARIABLE(0);
Requirements req("CUDNN AVGPOOL3d OP");
req.expectEq(makeInfoVariable(input->dataType(), TYPE_MSG_INPUT),
makeInfoVariable(output->dataType(), TYPE_MSG_OUTPUT)) &&
req.expectIn(makeInfoVariable(input->dataType(), TYPE_MSG_INPUT),
{INT32, HALF, FLOAT32, DOUBLE});
req.logTheSuccess();
return req;
}
//////////////////////////////////////////////////////////////////////////
PLATFORM_IMPL(avgpool3dnew_bp, ENGINE_CUDA) {
auto input = INPUT_VARIABLE(0); // [bS, iD, iH, iW, iC] (NDHWC) or [bS, iC, iD, iH, iW] (NCDHW)
auto gradO = INPUT_VARIABLE(1); // [bS, oD, oH, oW, oC] (NDHWC) or [bS, oC, oD, oH, oW] (NCDHW), epsilon_next
auto gradI = OUTPUT_VARIABLE(0); // [bS, iD, iH, iW, iC] (NDHWC) or [bS, iC, iD, iH, iW] (NCDHW), epsilon
const LongType kD = INT_ARG(0); // filter(kernel) depth
const LongType kH = INT_ARG(1); // filter(kernel) height
const LongType kW = INT_ARG(2); // filter(kernel) width
const LongType sD = INT_ARG(3); // strides depth
const LongType sH = INT_ARG(4); // strides height
const LongType sW = INT_ARG(5); // strides width
LongType pD = INT_ARG(6); // paddings depth
LongType pH = INT_ARG(7); // paddings height
LongType pW = INT_ARG(8); // paddings width
const LongType dD = INT_ARG(9); // dilations depth
const LongType dH = INT_ARG(10); // dilations height
const LongType dW = INT_ARG(11); // dilations width
const int isSameMode = INT_ARG(12); // 1-SAME, 0-VALID
const int extraParam0 = INT_ARG(13); // define what divisor to use while averaging
const int isNCDHW = block.getIArguments()->size() > 14 ? !INT_ARG(14) : 1; // 0-NCDHW, 1-NDHWC
REQUIRE_TRUE(input->rankOf() == 5, 0, "AVGPOOL3DNEW_BP CUDNN OP: input should have rank of 5, but got %i instead",
input->rankOf());
REQUIRE_TRUE(dD != 0 && dH != 0 && dW != 0, 0,
"AVGPOOL3DNEW_BP CUDNN OP: dilation must not be zero, but got instead {%i, %i, %i}", dD, dH, dW);
LongType bS, iC, iD, iH, iW, oC, oD, oH,
oW; // batch size, input channels, input depth/height/width, output channels, output depth/height/width;
LongType indIOioC, indIOioD, indWoC, indWiC, indWkD; // corresponding indexes
ConvolutionUtils::getSizesAndIndexesConv3d(isNCDHW, 0, *input, *gradO, bS, iC, iD, iH, iW, oC, oD, oH, oW, indIOioC,
indIOioD, indWiC, indWoC, indWkD);
std::vector<LongType> expectedGradOShape =
ShapeUtils::composeShapeUsingDimsAndIdx({bS, iC, oD, oH, oW, 0, indIOioC, indIOioD, indIOioD + 1, indIOioD + 2});
std::vector<LongType> expectedGradIShape =
ShapeUtils::composeShapeUsingDimsAndIdx({bS, iC, iD, iH, iW, 0, indIOioC, indIOioD, indIOioD + 1, indIOioD + 2});
REQUIRE_TRUE(gradO->isSameShape(expectedGradOShape), 0,
"AVGPOOL3DNEW_BP CUDNN: wrong shape of output's gradients array (next epsilon), expected is %s, but got "
"%s instead !",
ShapeUtils::shapeAsString(expectedGradOShape).c_str(), ShapeUtils::shapeAsString(gradO).c_str());
REQUIRE_TRUE(
gradI->isSameShape(expectedGradIShape), 0,
"AVGPOOL3DNEW_BP CUDNN: wrong shape of input's gradients array (epsilon), expected is %s, but got %s instead !",
ShapeUtils::shapeAsString(expectedGradIShape).c_str(), ShapeUtils::shapeAsString(gradI).c_str());
if (isSameMode) // SAME
ConvolutionUtils::calcPadding3D(pD, pH, pW, oD, oH, oW, iD, iH, iW, kD, kH, kW, sD, sH, sW, dD, dH, dW);
const cudnnPoolingMode_t mode =
(extraParam0 == 0) ? CUDNN_POOLING_AVERAGE_COUNT_EXCLUDE_PADDING : CUDNN_POOLING_AVERAGE_COUNT_INCLUDE_PADDING;
pooling3dBpCUDNN(block.launchContext(), input, gradO, gradI, kD, kH, kW, sD, sH, sW, pD, pH, pW, dD, dH, dW, isNCDHW,
mode);
return Status::OK;
}
PLATFORM_CHECK(avgpool3dnew_bp, ENGINE_CUDA) {
auto input = INPUT_VARIABLE(0); // [bS, iD, iH, iW, iC] (NDHWC) or [bS, iC, iD, iH, iW] (NCDHW)
auto gradO = INPUT_VARIABLE(1); // [bS, oD, oH, oW, oC] (NDHWC) or [bS, oC, oD, oH, oW] (NCDHW), epsilon_next
auto gradI = OUTPUT_VARIABLE(0); // [bS, iD, iH, iW, iC] (NDHWC) or [bS, iC, iD, iH, iW] (NCDHW), epsilon
Requirements req("CUDNN AVGPOOL3d_BP OP");
req.expectEq(makeInfoVariable(input->dataType(), TYPE_MSG_INPUT0),
makeInfoVariable(gradO->dataType(), TYPE_MSG_INPUT1)) &&
req.expectEq(makeInfoVariable(input->dataType(), TYPE_MSG_INPUT),
makeInfoVariable(gradI->dataType(), TYPE_MSG_OUTPUT)) &&
req.expectIn(makeInfoVariable(input->dataType(), TYPE_MSG_INPUT),
{INT32, HALF, FLOAT32, DOUBLE}) &&
req.expect(
makeShapeInfoVariable(input, SHAPE_MSG_INPUT0), makeShapeInfoVariable(gradI, SHAPE_MSG_OUTPUT),
[](const decltype(input)& l, const decltype(gradI)& r) {
return shape::haveSameShapeAndStrides(l->shapeInfo(), r->shapeInfo());
},
EXPECTED_EQ_MSG);
req.logTheSuccess();
return req;
}
} // namespace platforms
} // namespace ops
} // namespace sd
libnd4j/include/ops/declarable/platform/cudnn/batchnorm.cu
+557
View File
@@ -0,0 +1,557 @@
/* ******************************************************************************
*
*
* This program and the accompanying materials are made available under the
* terms of the Apache License, Version 2.0 which is available at
* https://www.apache.org/licenses/LICENSE-2.0.
*
* See the NOTICE file distributed with this work for additional
* information regarding copyright ownership.
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
* WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
* License for the specific language governing permissions and limitations
* under the License.
*
* SPDX-License-Identifier: Apache-2.0
******************************************************************************/
//
// @author Yurii Shyrma (iuriish@yahoo.com)
//
#include <ops/declarable/helpers/convolutions.h>
#include "cudnnUtils.h"
namespace sd {
namespace ops {
namespace platforms {
//////////////////////////////////////////////////////////////////////////
static void batchnormCUDNN(const LaunchContext* context, NDArray* input, NDArray* mean,
NDArray* variance, NDArray* gamma, NDArray* beta, NDArray* output,
const double epsilon, const bool isSpatialMode) {
// input, output -> 4D:nchw, 5D:ncdhw
// mean, variance, gamma, beta -> 1xCx1x1 for 4D and 1xCx1x1x1 for 5D for BATCHNORM_MODE_SPATIAL mode
// -> 1xCxHxW for 4D and 1xCxDxHxW for 5D for BATCHNORM_MODE_PER_ACTIVATION mode
const cudnnDataType_t dataType = cudnnDataType(input->dataType());
const LongType xRank = input->rankOf();
auto handle = reinterpret_cast<cudnnHandle_t*>(context->getCuDnnHandle());
CHECK_CUDNN_FAILURE(cudnnSetStream(*handle, *context->getCudaStream()));
const std::vector<int> xShape = input->getShapeAsVectorInt(); // input and output have same shapes
std::vector<int> paramsShape, paramsStrides; // mean, variance, gamma and beta have same shapes
if (isSpatialMode) { // 1xCx1x1
const int iC = static_cast<int>(mean->lengthOf());
const int stride0 = static_cast<int>(mean->strideAt(0));
paramsShape = xRank == 4 ? std::vector<int>({1, iC, 1, 1}) : std::vector<int>({1, iC, 1, 1, 1});
paramsStrides = xRank == 4 ? std::vector<int>({iC * stride0, stride0, 1, 1})
: std::vector<int>({iC * stride0, stride0, 1, 1, 1});
} else {
paramsShape = std::vector<int>(mean->getShapeAsVector().begin(), mean->getShapeAsVector().end());
paramsStrides = xRank == 4
? std::vector<int>({static_cast<int>(mean->strideAt(0)), static_cast<int>(mean->strideAt(1)), static_cast<int>(mean->strideAt(2)),
static_cast<int>(mean->strideAt(3))})
: std::vector<int>({static_cast<int>(mean->strideAt(0)), static_cast<int>(mean->strideAt(1)), static_cast<int>(mean->strideAt(2)),
static_cast<int>(mean->strideAt(3)), static_cast<int>(mean->strideAt(4))});
}
std::vector<int> xStrides = {static_cast<int>(input->strideAt(0)), static_cast<int>(input->strideAt(1)), static_cast<int>(input->strideAt(2)),
static_cast<int>(input->strideAt(3))};
std::vector<int> zStrides = {static_cast<int>(output->strideAt(0)), static_cast<int>(output->strideAt(1)), static_cast<int>(output->strideAt(2)),
static_cast<int>(output->strideAt(3))};
if (xRank > 4) { // 5D
xStrides.push_back((LongType)input->strideAt(4));
zStrides.push_back((LongType)output->strideAt(4));
}
cudnnTensorFormat_t format = CUDNN_TENSOR_NCHW;
// input descriptor
x.set(dataType, xRank, xShape.data(), xStrides.data());
// output descriptor
CudnnTensor z;
z.set(dataType, xRank, xShape.data(), zStrides.data());
// mean, variance, gamma and beta descriptor, the same descriptor for all of them
CudnnTensor params;
params.set(dataType, xRank, paramsShape.data(), paramsStrides.data());
// provide scaling parameters
const float alpha32(1), beta32(0);
const double alpha64(1), beta64(0);
const void* ptrAlpha =
output->sizeOfT() <= 4 ? reinterpret_cast<const void*>(&alpha32) : reinterpret_cast<const void*>(&alpha64);
const void* ptrBeta =
output->sizeOfT() <= 4 ? reinterpret_cast<const void*>(&beta32) : reinterpret_cast<const void*>(&beta64);
NDArray::prepareSpecialUse({output}, {input, mean, variance, gamma, beta});
// calculations
CHECK_CUDNN_FAILURE_MSG(
STRINGIZE(cudnnBatchNormalizationForwardInference),
cudnnBatchNormalizationForwardInference(
*handle, isSpatialMode ? CUDNN_BATCHNORM_SPATIAL : CUDNN_BATCHNORM_PER_ACTIVATION, ptrAlpha, ptrBeta, x,
input->specialBuffer(), z, output->specialBuffer(), params, gamma->specialBuffer(), beta->specialBuffer(),
mean->specialBuffer(), variance->specialBuffer(), epsilon));
auto cudaErr = cudaStreamSynchronize(*context->getCudaStream());
if (cudaErr != 0) throw cuda_exception::build("batchnormCUDNN: cudaStreamSynchronize failed !", cudaErr);
NDArray::registerSpecialUse({output}, {input, mean, variance, gamma, beta});
}
//////////////////////////////////////////////////////////////////////////
static void batchnormBpCUDNN(const LaunchContext* context, NDArray* input, NDArray* mean,
NDArray* variance, NDArray* gamma, NDArray* gradO, NDArray* gradI,
NDArray* gradG, NDArray* gradB, const double epsilon, const bool isSpatialMode) {
// input, gradO, gradI -> 4D:nchw, 5D:ncdhw
// mean, variance, gamma, beta, gradM, gradV, gradG, gradB -> 1xCx1x1 for 4D and 1xCx1x1x1 for 5D for
// BATCHNORM_MODE_SPATIAL mode
// -> 1xCxHxW for 4D and 1xCxDxHxW for 5D for
// BATCHNORM_MODE_PER_ACTIVATION mode
const cudnnDataType_t dataType = cudnnDataType(input->dataType());
const int xRank = input->rankOf();
auto handle = reinterpret_cast<cudnnHandle_t*>(context->getCuDnnHandle());
cudnnStatus_t err = cudnnSetStream(*handle, *context->getCudaStream());
const std::vector<int> xShape = input->getShapeAsVectorInt(); // input and output have same shapes
std::vector<int> paramsShape, paramsStrides; // mean, variance, gamma and beta have same shapes
if (isSpatialMode) { // 1xCx1x1
const int iC = static_cast<int>(mean->lengthOf());
const int stride0 = static_cast<int>(mean->strideAt(0));
paramsShape = xRank == 4 ? std::vector<int>({1, iC, 1, 1}) : std::vector<int>({1, iC, 1, 1, 1});
paramsStrides = xRank == 4 ? std::vector<int>({iC * stride0, stride0, 1, 1})
: std::vector<int>({iC * stride0, stride0, 1, 1, 1});
} else {
paramsShape = std::vector<int>(mean->getShapeAsVector().begin(), mean->getShapeAsVector().end());
paramsStrides = xRank == 4
? std::vector<int>({static_cast<int>(mean->strideAt(0)), static_cast<int>(mean->strideAt(1)), static_cast<int>(mean->strideAt(2)),
static_cast<int>(mean->strideAt(3))})
: std::vector<int>({static_cast<int>(mean->strideAt(0)), static_cast<int>(mean->strideAt(1)), static_cast<int>(mean->strideAt(2)),
static_cast<int>(mean->strideAt(3)), static_cast<int>(mean->strideAt(4))});
}
std::vector<int> xStrides = {static_cast<int>(input->strideAt(0)), static_cast<int>(input->strideAt(1)), static_cast<int>(input->strideAt(2)),
static_cast<int>(input->strideAt(3))};
std::vector<int> dxStrides = {static_cast<int>(gradI->strideAt(0)), static_cast<int>(gradI->strideAt(1)), static_cast<int>(gradI->strideAt(2)),
static_cast<int>(gradI->strideAt(3))};
std::vector<int> dzStrides = {static_cast<int>(gradO->strideAt(0)), static_cast<int>(gradO->strideAt(1)), static_cast<int>(gradO->strideAt(2)),
static_cast<int>(gradO->strideAt(3))};
if (xRank > 4) { // 5D
xStrides.push_back(static_cast<int>(input->strideAt(4)));
dxStrides.push_back(static_cast<int>(gradI->strideAt(4)));
dzStrides.push_back(static_cast<int>(gradO->strideAt(4)));
}
cudnnTensorFormat_t format = CUDNN_TENSOR_NCHW;
// input descriptor
CudnnTensor x;
x.set(dataType, xRank, xShape.data(), xStrides.data());
// gradO descriptor
CudnnTensor dz;
dz.set(dataType, xRank, xShape.data(), dzStrides.data());
// gradI descriptor
CudnnTensor dx;
dx.set(dataType, xRank, xShape.data(), dxStrides.data());
// mean, variance, gamma, gradG and gradB descriptor, the same descriptor for all of them
CudnnTensor params;
params.set(dataType, xRank, paramsShape.data(), paramsStrides.data());
// provide scaling parameters
const float alpha32(1), beta32(0);
double alpha64(1), beta64(0);
const void* ptrAlpha =
input->sizeOfT() <= 4 ? reinterpret_cast<const void*>(&alpha32) : reinterpret_cast<const void*>(&alpha64);
const void* ptrBeta =
input->sizeOfT() <= 4 ? reinterpret_cast<const void*>(&beta32) : reinterpret_cast<const void*>(&beta64);
NDArray::prepareSpecialUse({gradI, gradG, gradB}, {input, mean, variance, gamma, gradO});
// calculations
// TODO: we can use cache here
CHECK_CUDNN_FAILURE_MSG(
STRINGIZE(cudnnBatchNormalizationBackward),
cudnnBatchNormalizationBackward(*handle, isSpatialMode ? CUDNN_BATCHNORM_SPATIAL : CUDNN_BATCHNORM_PER_ACTIVATION,
ptrAlpha, ptrBeta, ptrAlpha, ptrBeta, x, input->specialBuffer(), dz,
gradO->specialBuffer(), dx, gradI->specialBuffer(), params,
gamma->specialBuffer(), gradG->specialBuffer(), gradB->specialBuffer(), epsilon,
nullptr /*mean->specialBuffer()*/, nullptr /*variance->specialBuffer()*/));
auto cudaErr = cudaStreamSynchronize(*context->getCudaStream());
if (cudaErr != 0) throw cuda_exception::build("batchnormBpCUDNN: cudaStreamSynchronize failed !", cudaErr);
NDArray::registerSpecialUse({gradI, gradG, gradB}, {input, mean, variance, gamma, gradO});
}
//////////////////////////////////////////////////////////////////////////
PLATFORM_IMPL(batchnorm, ENGINE_CUDA) {
auto input = INPUT_VARIABLE(0);
auto mean = INPUT_VARIABLE(1);
auto variance = INPUT_VARIABLE(2);
NDArray* gamma = nullptr;
NDArray* beta = nullptr;
auto output = OUTPUT_VARIABLE(0);
const bool applyScale = (bool)INT_ARG(0);
const bool applyOffset = (bool)INT_ARG(1);
const double epsilon = T_ARG(0);
if (applyScale) gamma = INPUT_VARIABLE(3);
if (applyOffset) beta = INPUT_VARIABLE(3 + (int)applyScale);
const int numOfIntArgs = block.getIArguments()->size();
const int inRank = input->rankOf();
// get axes args to normalize input array over
std::vector<int> axes;
if (numOfIntArgs > 2)
for (int i = 2; i < numOfIntArgs; ++i) axes.push_back(INT_ARG(i));
else
axes.push_back(inRank - 1); // default dimension to reduce along is last dimension
const int numOfAxes = axes.size();
REQUIRE_TRUE(numOfAxes <= inRank, 0,
"BATCHNORM CUDNN op: too big number of input axes to normalize over, expected number should be less or "
"equal to rank of input array, but got %i and %i correspondingly !",
numOfAxes, inRank);
// evaluate expected shape for mean, variance and gamma. These 3 arrays should have identical shapes
// for example if input shape is {2,3,4,5,6} and axes = {1,3}, then expected shape would be {1,3,1,5,1}, and if axes =
// {3}, then expected shape would be {5}
std::vector<LongType> expShape;
if (numOfAxes == 1)
expShape.push_back(input->sizeAt(axes[0]));
else { // get, for example, something like {1, inputDim1, 1, inputDim3, 1} if axes = {1, 3}
expShape = std::vector<LongType>(inRank, 1);
for (LongType i = 0; i < numOfAxes; ++i) expShape[axes[i]] = input->sizeAt(axes[i]);
}
REQUIRE_TRUE(mean->isSameShape(expShape), 0,
"BATCHNORM CUDNN op: wrong shape of mean array, expected is %s, but got %s instead !",
ShapeUtils::shapeAsString(expShape).c_str(), ShapeUtils::shapeAsString(mean).c_str());
REQUIRE_TRUE(variance->isSameShape(expShape), 0,
"BATCHNORM CUDNN op: wrong shape of variance array, expected is %s, but got %s instead !",
ShapeUtils::shapeAsString(expShape).c_str(), ShapeUtils::shapeAsString(variance).c_str());
if (gamma)
REQUIRE_TRUE(gamma->isSameShape(expShape), 0,
"BATCHNORM CUDNN op: wrong shape of gamma array, expected is %s, but got %s instead !",
ShapeUtils::shapeAsString(expShape).c_str(), ShapeUtils::shapeAsString(gamma).c_str());
if (beta)
REQUIRE_TRUE(beta->isSameShape(expShape), 0,
"BATCHNORM CUDNN op: wrong shape of beta array, expected is %s, but got %s instead !",
ShapeUtils::shapeAsString(expShape).c_str(), ShapeUtils::shapeAsString(beta).c_str());
// types of all input arrays should be the same
for (int i = 1; i < block.width(); ++i)
REQUIRE_TRUE(INPUT_VARIABLE(0)->dataType() == INPUT_VARIABLE(i)->dataType(), 0,
"BATCHNORM CUDNN op: types of all input arrays should be the same !");
// cudnn supports NCHW format only
const bool needPermut = axes.size() == 1 && mean->lengthOf() == input->sizeAt(-1);
std::unique_ptr<NDArray> tmpGamma = {}, tmpBeta = {}, tmpInput = {}, tmpOutput = {};
if (needPermut) { // if NHWC
std::vector<LongType> perm =
inRank == 4 ? std::vector<LongType>({0, 3, 1, 2}) : std::vector<LongType>({0, 4, 1, 2, 3}); // NHWC -> NCHW
tmpInput.reset(new NDArray(input->permute(perm)));
tmpOutput.reset(new NDArray(output->permute(perm)));
input = tmpInput.get();
output = tmpOutput.get();
}
// cudnn requires gamma and beta to be non-nullptr
if (!applyScale) {
tmpGamma.reset(new NDArray(mean));
gamma = tmpGamma.get();
*gamma = 1;
}
if (!applyOffset) {
tmpBeta.reset(new NDArray(mean));
beta = tmpBeta.get();
*beta = 0;
}
// calculations
batchnormCUDNN(block.launchContext(), input, mean, variance, gamma, beta, output, epsilon, axes.size() == 1);
return Status::OK;
}
//////////////////////////////////////////////////////////////////////////
PLATFORM_CHECK(batchnorm, ENGINE_CUDA) {
const bool applyScale = (bool)INT_ARG(0);
const bool applyOffset = (bool)INT_ARG(1);
NDArray* input = INPUT_VARIABLE(0);
NDArray* mean = INPUT_VARIABLE(1);
NDArray* variance = INPUT_VARIABLE(2);
NDArray* gamma = applyScale ? INPUT_VARIABLE(3) : nullptr;
NDArray* beta = applyOffset ? INPUT_VARIABLE(3 + (int)applyScale) : nullptr;
const int numOfIntArgs = block.getIArguments()->size();
const int xRank = input->rankOf();
// *********************************** //
// get axes args to normalize input array over
std::vector<int> axes;
if (numOfIntArgs > 2)
for (int i = 2; i < numOfIntArgs; ++i) axes.push_back(INT_ARG(i));
else
axes.push_back(xRank - 1); // default dimension to reduce along is last dimension
Requirements req("CUDNN BATCHNORM OP");
req.expectIn(makeInfoVariable(xRank, RANK_MSG_INPUT0), {4, 5}) &&
req.expectIn(makeInfoVariable(input->dataType(), TYPE_MSG_INPUT0),
{HALF, FLOAT32, DOUBLE}) &&
req.expectIn(makeInfoVariable(axes.size(), "axes.size()"), {1, 3, 4}) &&
req.expect(
makeShapeInfoVariable(mean, SHAPE_MSG_INPUT1), makeShapeInfoVariable(variance, SHAPE_MSG_INPUT2),
[](const decltype(mean)& l, const decltype(variance)& r) {
return shape::haveSameShapeAndStrides(l->shapeInfo(), r->shapeInfo());
},
EXPECTED_EQ_MSG);
if (gamma) {
req.expect(
makeShapeInfoVariable(gamma, SHAPE_MSG_INPUT_ "#gamma"), makeShapeInfoVariable(mean, SHAPE_MSG_INPUT1),
[](const decltype(gamma)& l, const decltype(mean)& r) {
return shape::haveSameShapeAndStrides(l->shapeInfo(), r->shapeInfo());
},
EXPECTED_EQ_MSG);
}
if (beta) {
req.expect(
makeShapeInfoVariable(beta, SHAPE_MSG_INPUT_ "#beta"), makeShapeInfoVariable(mean, SHAPE_MSG_INPUT1),
[](const decltype(beta)& l, const decltype(mean)& r) {
return shape::haveSameShapeAndStrides(l->shapeInfo(), r->shapeInfo());
},
EXPECTED_EQ_MSG);
}
if (axes.size() == 1) {
req.expectIn(makeInfoVariable(mean->lengthOf(), LENGTH_MSG_INPUT1), {-1, 1});
} else {
auto inputShapeModif = input->getShapeAsVector(); // [dim0,dim1,dim2,dim3] 4D or [dim0,dim1,dim2,dim3,dim4]
inputShapeModif[0] = 1;
// mean [1,dim1,dim2,dim3] 4D or [1,dim1,dim2,dim3,dim4]
req.expect(
makeShapeInfoVariable(mean, SHAPE_MSG_INPUT1),
makeShapeInfoVariable(inputShapeModif, SHAPE_MSG_INPUT_ "#expect"),
[](const decltype(mean)& l, const decltype(inputShapeModif)& r) { return l->isSameShape(r); }, EXPECTED_EQ_MSG);
}
req.logTheSuccess();
return req;
}
//////////////////////////////////////////////////////////////////////////
PLATFORM_IMPL(batchnorm_bp, ENGINE_CUDA) {
NDArray* input = INPUT_VARIABLE(0);
NDArray* mean = INPUT_VARIABLE(1);
NDArray* variance = INPUT_VARIABLE(2);
NDArray* gamma = nullptr;
NDArray* beta = nullptr;
NDArray* gradO = INPUT_VARIABLE(block.width() - 1); // next epsilon
NDArray* gradI = OUTPUT_VARIABLE(0);
NDArray* gradM = OUTPUT_VARIABLE(1);
NDArray* gradV = OUTPUT_VARIABLE(2);
NDArray* gradG = nullptr;
NDArray* gradB = nullptr;
const bool applyScale = (bool)INT_ARG(0);
const bool applyOffset = (bool)INT_ARG(1);
const float epsilon = T_ARG(0);
if (applyScale) {
gamma = INPUT_VARIABLE(3);
gradG = OUTPUT_VARIABLE(3);
}
if (applyOffset) {
beta = INPUT_VARIABLE(3 + (int)applyScale);
gradB = OUTPUT_VARIABLE(3 + (int)applyScale);
}
const int numOfIntArgs = block.getIArguments()->size();
const int inRank = input->rankOf();
// get axes args to normalize input array over
std::vector<int> axes;
if (numOfIntArgs > 2)
for (int i = 2; i < numOfIntArgs; ++i) axes.push_back(INT_ARG(i));
else
axes.push_back(inRank - 1); // default dimension to reduce along is last dimension
const int numOfAxes = axes.size();
REQUIRE_TRUE(numOfAxes <= inRank, 0,
"BATCHNORM_BP CUDNN op: too big number of input axes to normalize over, expected number should be less "
"or equal to rank of input array, but got %i and %i correspondingly !",
numOfAxes, inRank);
// evaluate expected shape for mean, variance and gamma. These 3 arrays should have identical shapes
// for example if input shape is {2,3,4,5,6} and axes = {1,3}, then expected shape would be {1,3,1,5,1}, and if axes =
// {3}, then expected shape would be {5}
std::vector<LongType> expShape;
if (numOfAxes == 1)
expShape.push_back(input->sizeAt(axes[0]));
else { // get, for example, something like {1, inputDim1, 1, inputDim3, 1} if axes = {1, 3}
expShape = std::vector<LongType>(inRank, 1);
for (LongType i = 0; i < numOfAxes; ++i) expShape[axes[i]] = input->sizeAt(axes[i]);
}
REQUIRE_TRUE(mean->isSameShape(expShape), 0,
"BATCHNORM_BP CUDNN op: wrong shape of mean array, expected is %s, but got %s instead !",
ShapeUtils::shapeAsString(expShape).c_str(), ShapeUtils::shapeAsString(mean).c_str());
REQUIRE_TRUE(variance->isSameShape(expShape), 0,
"BATCHNORM_BP CUDNN op: wrong shape of variance array, expected is %s, but got %s instead !",
ShapeUtils::shapeAsString(expShape).c_str(), ShapeUtils::shapeAsString(variance).c_str());
if (gamma)
REQUIRE_TRUE(gamma->isSameShape(expShape), 0,
"BATCHNORM_BP CUDNN op: wrong shape of gamma array, expected is %s, but got %s instead !",
ShapeUtils::shapeAsString(expShape).c_str(), ShapeUtils::shapeAsString(gamma).c_str());
if (beta)
REQUIRE_TRUE(beta->isSameShape(expShape), 0,
"BATCHNORM_BP CUDNN op: wrong shape of beta array, expected is %s, but got %s instead !",
ShapeUtils::shapeAsString(expShape).c_str(), ShapeUtils::shapeAsString(beta).c_str());
REQUIRE_TRUE(input->isSameShape(gradO), 0,
"BATCHNORM_BP CUDNN op: wrong shape of output gradients array, expected is %s, but got %s instead !",
ShapeUtils::shapeAsString(input).c_str(), ShapeUtils::shapeAsString(gradO).c_str());
// types of all input arrays should be the same (except gradO)
for (int i = 1; i < block.width() - 2; ++i)
REQUIRE_TRUE(INPUT_VARIABLE(0)->dataType() == INPUT_VARIABLE(i)->dataType(), 0,
"BATCHNORM_BP CUDNN op: types of arrays (input, mean, variance, gamma, beta) should be the same !");
// cudnn supports NCHW format only
const bool needPermut = axes.size() == 1 && mean->lengthOf() != input->sizeAt(1);
std::unique_ptr<NDArray> tmpGamma = {}, tmpGradG = {}, tmpGradB = {}, tmpInput = {}, tmpGradI = {}, tmpGradO = {};
if (needPermut) { // if NHWC
std::vector<LongType> perm =
inRank == 4 ? std::vector<LongType>({0, 3, 1, 2}) : std::vector<LongType>({0, 4, 1, 2, 3}); // NHWC -> NCHW
tmpInput.reset(new NDArray(input->permute(perm)));
tmpGradO.reset(new NDArray(gradO->permute(perm)));
tmpGradI.reset(new NDArray(gradI->permute(perm)));
input = tmpInput.get();
gradO = tmpGradO.get();
gradI = tmpGradI.get();
}
// cudnn requires gamma, gradG, gradB to be non-nullptr
if (!applyScale) {
tmpGamma.reset(new NDArray(mean));
tmpGradG.reset(new NDArray(mean));
gamma = tmpGamma.get();
gradG = tmpGradG.get();
*gamma = 1;
}
if (!applyOffset) {
tmpGradB.reset(new NDArray(mean));
gradB = tmpGradB.get();
}
// calculations
batchnormBpCUDNN(block.launchContext(), input, mean, variance, gamma, gradO, gradI, gradG, gradB, epsilon,
axes.size() == 1);
*gradM = 0; // put zeros so far
*gradV = 0; // put zeros so far
return Status::OK;
}
PLATFORM_CHECK(batchnorm_bp, ENGINE_CUDA) {
NDArray* input = INPUT_VARIABLE(0);
NDArray* mean = INPUT_VARIABLE(1);
NDArray* variance = INPUT_VARIABLE(2);
NDArray* gamma = nullptr;
NDArray* beta = nullptr;
NDArray* gradO = INPUT_VARIABLE(block.width() - 1); // next epsilon
NDArray* gradI = OUTPUT_VARIABLE(0);
NDArray* gradM = OUTPUT_VARIABLE(1);
NDArray* gradV = OUTPUT_VARIABLE(2);
NDArray* gradG = nullptr;
NDArray* gradB = nullptr;
const int numOfIntArgs = block.getIArguments()->size();
const int xRank = input->rankOf();
// *********************************** //
// get axes args to normalize input array over
std::vector<int> axes;
if (numOfIntArgs > 2)
for (int i = 2; i < numOfIntArgs; ++i) axes.push_back(INT_ARG(i));
else
axes.push_back(xRank - 1); // default dimension to reduce along is last dimension
Requirements req("CUDNN BATCHNORM_BP OP");
req.expectIn(makeInfoVariable(xRank, RANK_MSG_INPUT0), {4, 5}) &&
req.expectIn(makeInfoVariable(input->dataType(), TYPE_MSG_INPUT0),
{HALF, FLOAT32, DOUBLE}) &&
req.expectIn(makeInfoVariable(axes.size(), "axes.size()"), {1, 3, 4}) &&
req.expect(
makeShapeInfoVariable(mean, SHAPE_MSG_INPUT1), makeShapeInfoVariable(variance, SHAPE_MSG_INPUT2),
[](const decltype(mean)& l, const decltype(variance)& r) {
return shape::haveSameShapeAndStrides(l->shapeInfo(), r->shapeInfo());
},
EXPECTED_EQ_MSG);
if (gamma) {
req.expect(
makeShapeInfoVariable(gamma, SHAPE_MSG_INPUT_ "#gamma"), makeShapeInfoVariable(mean, SHAPE_MSG_INPUT1),
[](const decltype(gamma)& l, const decltype(mean)& r) {
return shape::haveSameShapeAndStrides(l->shapeInfo(), r->shapeInfo());
},
EXPECTED_EQ_MSG);
}
if (gradG) {
req.expect(
makeShapeInfoVariable(gradG, SHAPE_MSG_INPUT_ "#gradG"), makeShapeInfoVariable(mean, SHAPE_MSG_INPUT1),
[](const decltype(gradG)& l, const decltype(mean)& r) {
return shape::haveSameShapeAndStrides(l->shapeInfo(), r->shapeInfo());
},
EXPECTED_EQ_MSG);
}
if (gradB) {
req.expect(
makeShapeInfoVariable(gradB, SHAPE_MSG_INPUT_ "#gradB"), makeShapeInfoVariable(mean, SHAPE_MSG_INPUT1),
[](const decltype(gradB)& l, const decltype(mean)& r) {
return shape::haveSameShapeAndStrides(l->shapeInfo(), r->shapeInfo());
},
EXPECTED_EQ_MSG);
}
if (axes.size() == 1) {
// isFormatGood = mean->lengthOf() == input->sizeAt(1) || mean->lengthOf() == input->sizeAt(-1); // mean [C]
req.expectIn(makeInfoVariable(mean->lengthOf(), LENGTH_MSG_INPUT1), {-1, 1});
} else {
auto inputShapeModif = input->getShapeAsVector(); // [dim0,dim1,dim2,dim3] 4D or [dim0,dim1,dim2,dim3,dim4]
inputShapeModif[0] = 1;
// isFormatGood = mean->isSameShape(inputShapeModif); // mean [1,dim1,dim2,dim3] 4D or
// [1,dim1,dim2,dim3,dim4]
req.expect(
makeShapeInfoVariable(mean, SHAPE_MSG_INPUT1),
makeShapeInfoVariable(inputShapeModif, SHAPE_MSG_INPUT_ "#expect"),
[](const decltype(mean)& l, const decltype(inputShapeModif)& r) { return l->isSameShape(r); }, EXPECTED_EQ_MSG);
}
req.logTheSuccess();
return req;
}
} // namespace platforms
} // namespace ops
} // namespace sd
libnd4j/include/ops/declarable/platform/cudnn/conv2d.cu
+529
View File
@@ -0,0 +1,529 @@
/* ******************************************************************************
*
*
* This program and the accompanying materials are made available under the
* terms of the Apache License, Version 2.0 which is available at
* https://www.apache.org/licenses/LICENSE-2.0.
*
* See the NOTICE file distributed with this work for additional
* information regarding copyright ownership.
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
* WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
* License for the specific language governing permissions and limitations
* under the License.
*
* SPDX-License-Identifier: Apache-2.0
******************************************************************************/
//
// @author raver119@gmail.com
// @author Yurii Shyrma (iuriish@yahoo.com)
//
#include <ops/declarable/helpers/convolutions.h>
#include "cudnnUtils.h"
namespace sd {
namespace ops {
namespace platforms {
//////////////////////////////////////////////////////////////////////////
static void conv2dCUDNN(const LaunchContext* context, NDArray* input, NDArray* weights, NDArray* bias,
NDArray* output, const int kH, const LongType kW, const LongType sH, const LongType sW, const LongType pH,
const LongType pW, const LongType dH, const LongType dW, const int paddingMode, const bool isNCHW,
const int wFormat) {
// cudnn support only two formats for weights {oC,iC,kH,kW} and {oC,kH,kW,iC}
LongType bS, iC, iH, iW, oC, oH,
oW; // batch size, input channels, input height/width, output channels, output height/width;
LongType indIOioC, indIiH, indWoC, indWiC, indWkH, indOoH; // corresponding indexes
ConvolutionUtils::getSizesAndIndexesConv2d(isNCHW, wFormat, *input, *output, bS, iC, iH, iW, oC, oH, oW, indIOioC,
indIiH, indWiC, indWoC, indWkH, indOoH);
auto handle = reinterpret_cast<cudnnHandle_t*>(context->getCuDnnHandle());
CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnSetStream), cudnnSetStream(*handle, *context->getCudaStream()));
cudnnTensorFormat_t format = isNCHW ? CUDNN_TENSOR_NCHW : CUDNN_TENSOR_NHWC;
cudnnTensorFormat_t formatW = 0 == wFormat ? format : (1 == wFormat ? CUDNN_TENSOR_NCHW : CUDNN_TENSOR_NHWC);
// input descriptor
CudnnTensor x;
if (input->ordering() == 'c')
x.set4D(format, cudnnDataType(input->dataType()), bS, iC, iH, iW);
else
x.set4DEx(cudnnDataType(input->dataType()), bS, iC, iH, iW, input->strideAt(0), input->strideAt(indIOioC),
input->strideAt(indIiH), input->strideAt(indIiH + 1));
// weights descriptor
FilterDesc w;
w.set4D(cudnnDataType(weights->dataType()), formatW, oC, iC, kH, kW);
// output descriptor
CudnnTensor z;
if (output->ordering() == 'c')
z.set4D(format, cudnnDataType(output->dataType()), bS, oC, oH, oW);
else
z.set4DEx(cudnnDataType(output->dataType()), bS, oC, oH, oW, output->strideAt(0), output->strideAt(indIOioC),
output->strideAt(indOoH), output->strideAt(indOoH + 1));
// description of convolution
ConvolutionDesc conv;
conv.set2D(pH, pW, sH, sW, dH, dW, CUDNN_CROSS_CORRELATION, cudnnDataType(output->dataType()));
// algorithm description
cudnnConvolutionFwdAlgo_t algo;
cudnnConvolutionFwdAlgoPerf_t algoPerf;
int count = 0;
// err = cudnnGetConvolutionForwardAlgorithm(*handle, x, w, conv, z, CUDNN_CONVOLUTION_FWD_PREFER_FASTEST, 0, &algo);
CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnFindConvolutionForwardAlgorithm),
cudnnFindConvolutionForwardAlgorithm(*handle, x, w, conv, z, 1, &count, &algoPerf));
if (count == 0)
throw cuda_exception::build("conv2dCUDNN: cudnnGetConvolutionForwardAlgorithm failed as the count is 0", 0);
algo = algoPerf.algo;
PointersManager manager(context, __func__);
// allocate auxiliary device memory, abbreviation ws means workspace
size_t wsSize;
CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnGetConvolutionForwardWorkspaceSize),
cudnnGetConvolutionForwardWorkspaceSize(*handle, x, w, conv, z, algo, &wsSize));
void* wsData = manager.allocateDevMem(wsSize);
// provide scaling parameters
const float alpha32(1), beta32(0);
const double alpha64(1), beta64(0);
const void* alpha =
output->sizeOfT() <= 4 ? reinterpret_cast<const void*>(&alpha32) : reinterpret_cast<const void*>(&alpha64);
const void* beta =
output->sizeOfT() <= 4 ? reinterpret_cast<const void*>(&beta32) : reinterpret_cast<const void*>(&beta64);
NDArray::prepareSpecialUse({output}, {input, weights, bias});
// run calculation
CHECK_CUDNN_FAILURE_MSG(
STRINGIZE(cudnnConvolutionForward),
cudnnConvolutionForward(*handle, alpha, x, input->specialBuffer(), w, weights->specialBuffer(), conv, algo,
wsData, wsSize, beta, z, output->specialBuffer()));
// add bias if it is present
if (bias != nullptr) {
CudnnTensor b;
b.set4D(CUDNN_TENSOR_NCHW, cudnnDataType(bias->dataType()), 1, oC, 1, 1);
CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnAddTensor), cudnnAddTensor(*handle, alpha, b, bias->specialBuffer(), alpha,
z, output->specialBuffer()));
}
NDArray::registerSpecialUse({output}, {input, weights, bias});
}
//////////////////////////////////////////////////////////////////////////
static void conv2dBpCUDNN(const LaunchContext* context, NDArray* input, NDArray* weights,
NDArray* gradO, NDArray* gradI, NDArray* gradW, NDArray* gradB, const LongType kH,
const LongType kW, const LongType sH, const LongType sW, const LongType pH, const LongType pW, const LongType dH,
const LongType dW, const LongType paddingMode, const bool isNCHW, const int wFormat) {
LongType bS, iC, iH, iW, oC, oH,
oW; // batch size, input channels, input height/width, output channels, output height/width;
LongType indIOioC, indIiH, indWoC, indWiC, indWkH, indOoH; // corresponding indexes
ConvolutionUtils::getSizesAndIndexesConv2d(isNCHW, wFormat, *input, *gradO, bS, iC, iH, iW, oC, oH, oW, indIOioC,
indIiH, indWiC, indWoC, indWkH, indOoH);
auto handle = reinterpret_cast<cudnnHandle_t*>(context->getCuDnnHandle());
CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnSetStream), cudnnSetStream(*handle, *context->getCudaStream()));
cudnnTensorFormat_t format = isNCHW ? CUDNN_TENSOR_NCHW : CUDNN_TENSOR_NHWC;
cudnnTensorFormat_t formatW = 0 == wFormat ? format : (1 == wFormat ? CUDNN_TENSOR_NCHW : CUDNN_TENSOR_NHWC);
PointersManager manager(context, __func__);
// input descriptor, gradO descriptor, gradI descriptor
CudnnTensor x, dz, dx;
if (input->ordering() == 'c')
x.set4D(format, cudnnDataType(input->dataType()), bS, iC, iH, iW);
else
x.set4DEx(cudnnDataType(input->dataType()), bS, iC, iH, iW, input->strideAt(0), input->strideAt(indIOioC),
input->strideAt(indIiH), input->strideAt(indIiH + 1));
if (gradO->ordering() == 'c')
dz.set4D(format, cudnnDataType(gradO->dataType()), bS, oC, oH, oW);
else
dz.set4DEx(cudnnDataType(gradO->dataType()), bS, oC, oH, oW, gradO->strideAt(0), gradO->strideAt(indIOioC),
gradO->strideAt(indOoH), gradO->strideAt(indOoH + 1));
if (gradI->ordering() == 'c')
dx.set4D(format, cudnnDataType(gradI->dataType()), bS, iC, iH, iW);
else
dx.set4DEx(cudnnDataType(gradI->dataType()), bS, iC, iH, iW, gradI->strideAt(0), gradI->strideAt(indIOioC),
gradI->strideAt(indIiH), gradI->strideAt(indIiH + 1));
// gradW descriptor
FilterDesc dw;
dw.set4D(cudnnDataType(gradW->dataType()), formatW, oC, iC, kH, kW);
// description of convolution
ConvolutionDesc conv;
conv.set2D(pH, pW, sH, sW, dH, dW, CUDNN_CROSS_CORRELATION, cudnnDataType(gradO->dataType()));
// gradW algorithm description
cudnnConvolutionBwdFilterAlgo_t algoGradW;
cudnnConvolutionBwdFilterAlgoPerf_t algoGradWPerf;
int count = 0;
CHECK_CUDNN_FAILURE_MSG(
STRINGIZE(cudnnFindConvolutionBackwardFilterAlgorithm),
cudnnFindConvolutionBackwardFilterAlgorithm(*handle, x, dz, conv, dw, 1, &count, &algoGradWPerf));
if (count == 0)
throw cuda_exception::build(
"conv2dBpCUDNN: cudnnGetConvolutionBackwardFilterAlgorithm failed as the count is 0", 0);
algoGradW = algoGradWPerf.algo;
// gradI algorithm description
cudnnConvolutionBwdDataAlgo_t algoGradI;
cudnnConvolutionBwdDataAlgoPerf_t algoGradIPerf;
CHECK_CUDNN_FAILURE_MSG(
STRINGIZE(cudnnFindConvolutionBackwardDataAlgorithm),
cudnnFindConvolutionBackwardDataAlgorithm(*handle, dw, dz, conv, x, 1, &count, &algoGradIPerf));
if (count == 0)
throw cuda_exception::build("conv2dBpCUDNN: cudnnGetConvolutionBackwardDataAlgorithm failed as the count is 0",
0);
algoGradI = algoGradIPerf.algo;
// allocate auxiliary device memory for gradW calculation, abbreviation ws means workspace
size_t wsGradWSize;
CHECK_CUDNN_FAILURE_MSG(
STRINGIZE(cudnnGetConvolutionBackwardFilterWorkspaceSize),
cudnnGetConvolutionBackwardFilterWorkspaceSize(*handle, x, dz, conv, dw, algoGradW, &wsGradWSize));
void* wsGradWData = manager.allocateDevMem(wsGradWSize);
// allocate auxiliary device memory for gradI calculation, abbreviation ws means workspace
size_t wsGradISize;
CHECK_CUDNN_FAILURE_MSG(
STRINGIZE(cudnnGetConvolutionBackwardDataWorkspaceSize),
cudnnGetConvolutionBackwardDataWorkspaceSize(*handle, dw, dz, conv, dx, algoGradI, &wsGradISize));
void* wsGradIData = manager.allocateDevMem(wsGradISize);
// provide scaling parameters
const float alpha32(1), beta32(0);
const double alpha64(1), beta64(0);
const void* alpha =
gradO->sizeOfT() <= 4 ? reinterpret_cast<const void*>(&alpha32) : reinterpret_cast<const void*>(&alpha64);
const void* beta =
gradO->sizeOfT() <= 4 ? reinterpret_cast<const void*>(&beta32) : reinterpret_cast<const void*>(&beta64);
NDArray::prepareSpecialUse({gradI, gradW, gradB}, {input, weights, gradO});
// run calculation for gradB (if not nullptr)
if (gradB != nullptr) {
CudnnTensor db;
db.set4D(CUDNN_TENSOR_NCHW, cudnnDataType(gradB->dataType()), 1, oC, 1, 1);
CHECK_CUDNN_FAILURE_MSG(
STRINGIZE(cudnnConvolutionBackwardBias),
cudnnConvolutionBackwardBias(*handle, alpha, dz, gradO->specialBuffer(), beta, db, gradB->specialBuffer()));
}
// run calculation for gradW
CHECK_CUDNN_FAILURE_MSG(
STRINGIZE(cudnnConvolutionBackwardFilter),
cudnnConvolutionBackwardFilter(*handle, alpha, x, input->specialBuffer(), dz, gradO->specialBuffer(), conv,
algoGradW, wsGradWData, wsGradWSize, beta, dw, gradW->specialBuffer()));
// run calculation for gradI
CHECK_CUDNN_FAILURE_MSG(
STRINGIZE(cudnnConvolutionBackwardData),
cudnnConvolutionBackwardData(*handle, alpha, dw, weights->specialBuffer(), dz, gradO->specialBuffer(), conv,
algoGradI, wsGradIData, wsGradISize, beta, dx, gradI->specialBuffer()));
NDArray::registerSpecialUse({gradI, gradW, gradB}, {input, weights, gradO});
}
//////////////////////////////////////////////////////////////////////////
PLATFORM_IMPL(conv2d, ENGINE_CUDA) {
auto input = INPUT_VARIABLE(0); // [bS, iH, iW, iC] (NHWC) or [bS, iC, iH, iW] (NCHW)
auto weights = INPUT_VARIABLE(1); // [kH, kW, iC, oC], [oC, iC, kH, kW], [oC, kH, kW, iC]
auto bias = block.width() > 2 ? INPUT_VARIABLE(2) : nullptr; // [oC]
auto output = OUTPUT_VARIABLE(0); // [bS, oH, oW, oC] (NHWC) or [bS, oC, oH, oW] (NCHW)
LongType sH = INT_ARG(2); // strides height
LongType sW = INT_ARG(3); // strides width
LongType pH = INT_ARG(4); // paddings height
LongType pW = INT_ARG(5); // paddings width
LongType dH = INT_ARG(6); // dilations height
LongType dW = INT_ARG(7); // dilations width
int paddingMode = INT_ARG(8); // 0-VALID, 1-SAME
bool isNCHW = block.getIArguments()->size() > 9 ? !INT_ARG(9) : 1; // INT_ARG(9): 0-NCHW, 1-NHWC
int wFormat = block.getIArguments()->size() > 10
? INT_ARG(10)
: 0; // 0 - [kH, kW, iC, oC], 1 - [oC, iC, kH, kW], 2 - [oC, kH, kW, iC]
LongType kH = INT_ARG(0) > 0 ? INT_ARG(0) : static_cast<LongType>(weights->sizeAt(0)); // filter(kernel) height
LongType kW = INT_ARG(1) > 0 ? INT_ARG(1) : static_cast<LongType>(weights->sizeAt(1)); // filter(kernel) width
REQUIRE_TRUE(input->rankOf() == 4, 0,
"CUSTOM CONV2D CUDNN OP: rank of input array must be equal to 4, but got %i instead !", input->rankOf());
REQUIRE_TRUE(weights->rankOf() == 4, 0,
"CUSTOM CONV2D CUDNN OP: rank of weights array must be equal to 4, but got %i instead !",
weights->rankOf());
LongType bS, iC, iH, iW, oC, oH,
oW; // batch size, input channels, input height/width, output channels, output height/width;
LongType indIOioC, indIiH, indWoC, indWiC, indWkH, indOoH; // corresponding indexes
ConvolutionUtils::getSizesAndIndexesConv2d(isNCHW, wFormat, *input, *output, bS, iC, iH, iW, oC, oH, oW, indIOioC,
indIiH, indWiC, indWoC, indWkH, indOoH);
ConvolutionUtils::calcPadding2D(pH, pW, oH, oW, iH, iW, kH, kW, sH, sW, dH, dW, paddingMode);
std::vector<LongType> expectedWeightsShape = ConvolutionUtils::expectWeightsShape(wFormat, kH, kW, iC, oC);
REQUIRE_TRUE(weights->isSameShape(expectedWeightsShape), 0,
"CUSTOM CONV2D CUDNN OP: wrong shape of weights array, expected is %s, but got %s instead !",
ShapeUtils::shapeAsString(expectedWeightsShape).c_str(), ShapeUtils::shapeAsString(weights).c_str());
if (bias) {
REQUIRE_TRUE(bias->rankOf() <= 2 && oC == bias->lengthOf(), 0,
"CUSTOM CONV2D CUDNN OP: wrong shape of array with biases, expected rank, length: <=2, %i, but got "
"%i, %i instead !",
oC, bias->rankOf(), bias->lengthOf());
REQUIRE_TRUE((bias->rankOf() == 1 && bias->strideAt(0) == 1) ||
(bias->rankOf() == 2 && bias->sizeAt(0) == 1 && bias->strideAt(1) == 1) ||
(bias->rankOf() == 2 && bias->sizeAt(1) == 1 && bias->strideAt(0) == 1),
0, "CUSTOM CONV2D CUDNN OP: bias array should be contiguous in memory !");
}
std::unique_ptr<NDArray> tmpWeight = {}, tmpInput = {};
NDArray* newWeights = weights; // cudnn support only two formats {oC,iC,kH,kW} and {oC,kH,kW,iC}
if (0 == wFormat) {
// Create named vectors as lvalues
std::vector<LongType> nchwShape = {oC, iC, kH, kW};
std::vector<LongType> nhwcShape = {oC, kH, kW, iC};
// Use the appropriate one for the weight reset
tmpWeight.reset(
new NDArray(weights->ordering(),
isNCHW ? nchwShape : nhwcShape,
weights->dataType(), weights->getContext()));
newWeights = tmpWeight.get();
// Create named vectors as lvalues
std::vector<LongType> nchwDims = {3, 2, 0, 1};
std::vector<LongType> nhwcDims = {3, 0, 1, 2};
// Use the appropriate one in the call
NDArray assign = weights->permute(
isNCHW ? nchwDims : nhwcDims,
true, // copyToNewBuff
true); // resetStrides
newWeights->assign(&assign); // (kH, kW, iC, oC --> oC, iC, kH, kW) or (kH, kW, iC, oC --> oC, kH, kW, iC)
}
if (paddingMode == 1) { // in same paddingMode cudnn doesn't support asymmetric left/right top/bottopm paddings
auto ret = checkConv2dCUDNNPadAsymmetric(input, nullptr, iH, iW, oH, oW, kH, kW, sH, sW, pH, pW, dH, dW, isNCHW);
tmpInput = std::move(std::get<0>(ret)); // prolong life
if (tmpInput) input = tmpInput.get();
}
conv2dCUDNN(block.launchContext(), input, newWeights, bias, output, kH, kW, sH, sW, pH, pW, dH, dW, paddingMode,
isNCHW, wFormat);
return Status::OK;
}
//////////////////////////////////////////////////////////////////////////
PLATFORM_CHECK(conv2d, ENGINE_CUDA) {
auto input = INPUT_VARIABLE(0); // [bS, iH, iW, iC] (NHWC) or [bS, iC, iH, iW] (NCHW)
auto weights = INPUT_VARIABLE(1); // [kH, kW, iC, oC] always
auto bias = block.width() > 2 ? INPUT_VARIABLE(2) : nullptr; // [oC]
const int paddingMode = INT_ARG(8); // 0-VALID, 1-SAME, 2-CAUSAL
const bool badInputType = input->dataType() != DOUBLE && input->dataType() != FLOAT32 &&
input->dataType() != HALF;
const bool badWeightsType = weights->dataType() != DOUBLE && weights->dataType() != FLOAT32 &&
weights->dataType() != HALF;
const bool badBiasType = bias == nullptr
? false
: (bias->dataType() != DOUBLE && bias->dataType() != FLOAT32 &&
bias->dataType() != HALF);
return paddingMode != 2 && !badInputType && !badWeightsType && !badBiasType;
Requirements req("CUDNN CONV2d OP");
req.expectNotEq(makeInfoVariable(paddingMode, "paddingMode"), 2) &&
req.expectIn(makeInfoVariable(input->dataType(), TYPE_MSG_INPUT0),
{HALF, FLOAT32, DOUBLE}) &&
req.expectIn(makeInfoVariable(weights->dataType(), TYPE_MSG_INPUT1),
{HALF, FLOAT32, DOUBLE});
if (bias) {
req.expectIn(makeInfoVariable(bias->dataType(), TYPE_MSG_INPUT_ "#bias"),
{HALF, FLOAT32, DOUBLE});
}
req.logTheSuccess();
return req;
}
//////////////////////////////////////////////////////////////////////////
PLATFORM_IMPL(conv2d_bp, ENGINE_CUDA) {
auto input = INPUT_VARIABLE(0); // [bS, iH, iW, iC] (NHWC) or [bS, iC, iH, iW] (NCHW)
auto weights = INPUT_VARIABLE(1); // [kH, kW, iC, oC], [oC, iC, kH, kW], [oC, kH, kW, iC]
auto bias = block.width() > 3 ? INPUT_VARIABLE(2) : nullptr; // [oC]
auto gradO = block.width() > 3
? INPUT_VARIABLE(3)
: INPUT_VARIABLE(2); // [bS, oH, oW, oC] (NHWC) or [bS, oC, oH, oW] (NCHW), epsilon_next
auto gradI = OUTPUT_VARIABLE(0); // [bS, iH, iW, iC] (NHWC) or [bS, iC, iH, iW] (NCHW), epsilon
auto gradW = OUTPUT_VARIABLE(1); // [kH, kW, iC, oC], [oC, iC, kH, kW], [oC, kH, kW, iC]
auto gradB = block.width() > 3 ? OUTPUT_VARIABLE(2) : nullptr; // [oC]
LongType kH = INT_ARG(0); // filter(kernel) height
LongType kW = INT_ARG(1); // filter(kernel) width
LongType sH = INT_ARG(2); // strides height
LongType sW = INT_ARG(3); // strides width
LongType pH = INT_ARG(4); // paddings height
LongType pW = INT_ARG(5); // paddings width
LongType dH = INT_ARG(6); // dilations height
LongType dW = INT_ARG(7); // dilations width
int paddingMode = INT_ARG(8); // 0-VALID, 1-SAME
int isNCHW = block.getIArguments()->size() > 9 ? !INT_ARG(9) : 1; // INT_ARG(9): 0-NCHW, 1-NHWC
int wFormat = block.getIArguments()->size() > 10
? INT_ARG(10)
: 0; // 0 - [kH, kW, iC, oC], 1 - [oC, iC, kH, kW], 2 - [oC, kH, kW, iC]
REQUIRE_TRUE(input->rankOf() == 4, 0,
"CUSTOM CONV2D_BP CUDNN OP: rank of input array must be equal to 4, but got %i instead !",
input->rankOf());
REQUIRE_TRUE(weights->rankOf() == 4, 0,
"CUSTOM CONV2D_BP CUDNN OP: rank of weights array must be equal to 4, but got %i instead !",
weights->rankOf());
REQUIRE_TRUE(gradO->rankOf() == 4, 0,
"CUSTOM CONV2D_BP CUDNN OP: rank of output's gradients (next epsilon) array must be equal to 4, but got "
"%i instead !",
gradO->rankOf());
LongType bS, iC, iH, iW, oC, oH,
oW; // batch size, input channels, input height/width, output channels, output height/width;
LongType indIOioC, indIiH, indWoC, indWiC, indWkH, indOoH; // corresponding indexes
ConvolutionUtils::getSizesAndIndexesConv2d(isNCHW, wFormat, *input, *gradO, bS, iC, iH, iW, oC, oH, oW, indIOioC,
indIiH, indWiC, indWoC, indWkH, indOoH);
LongType trueoH, trueoW; // true output height, width
ConvolutionUtils::calcOutSizePool2D(trueoH, trueoW, kH, kW, sH, sW, pH, pW, dH, dW, iH, iW, paddingMode);
ConvolutionUtils::calcPadding2D(pH, pW, oH, oW, iH, iW, kH, kW, sH, sW, dH, dW, paddingMode);
std::vector<LongType> expectedGradOShape =
ShapeUtils::composeShapeUsingDimsAndIdx({bS, oC, trueoH, trueoW, 0, indIOioC, indOoH, indOoH + 1});
std::vector<LongType> expectedWeightsShape = ConvolutionUtils::expectWeightsShape(wFormat, kH, kW, iC, oC);
REQUIRE_TRUE(gradO->isSameShape(expectedGradOShape), 0,
"CUSTOM CONV2D_BP CUDNN OP: wrong shape of output gradients (next epsilon) array, expected is %s, but "
"got %s instead !",
ShapeUtils::shapeAsString(expectedGradOShape).c_str(), ShapeUtils::shapeAsString(gradO).c_str());
REQUIRE_TRUE(weights->isSameShape(expectedWeightsShape), 0,
"CUSTOM CONV2D_BP CUDNN OP: wrong shape of weights array, expected is %s, but got %s instead !",
ShapeUtils::shapeAsString(expectedWeightsShape).c_str(), ShapeUtils::shapeAsString(weights).c_str());
if (bias)
REQUIRE_TRUE(bias->rankOf() <= 2 && oC == bias->lengthOf(), 0,
"CUSTOM CONV2D_BP CUDNN OP: wrong shape of array with biases, expected rank, length: <=2, %i, but got "
"%i, %i instead !",
oC, bias->rankOf(), bias->lengthOf());
std::unique_ptr<NDArray> tmpGradI = {}, tmpInput = {}, tmpWeights = {}, tmpGradW = {};
NDArray *newWeights = weights, *newGradW = gradW; // cudnn support only two formats {oC,iC,kH,kW} and {oC,kH,kW,iC}
if (0 == wFormat) {
// Create named vectors as lvalues
std::vector<LongType> nchwGradShape = {oC, iC, kH, kW};
std::vector<LongType> nhwcGradShape = {oC, kH, kW, iC};
// Use the appropriate one for the gradW reset
tmpGradW.reset(
new NDArray(gradW->ordering(),
isNCHW ? nchwGradShape : nhwcGradShape,
gradW->dataType(), gradW->getContext()));
// Use the same vectors for the weights reset
tmpWeights.reset(
new NDArray(weights->ordering(),
isNCHW ? nchwGradShape : nhwcGradShape,
weights->dataType(), weights->getContext()));
newGradW = tmpGradW.get();
newWeights = tmpWeights.get();
// Create named vectors as lvalues
std::vector<LongType> nchwDims = {3, 2, 0, 1};
std::vector<LongType> nhwcDims = {3, 0, 1, 2};
NDArray assign = weights->permute(
isNCHW ? nchwDims : nhwcDims,
true, // copyToNewBuff
true);
// Use the appropriate one in the call
newWeights->assign(&assign);
}
NDArray* newInput = input;
NDArray* newGradI = gradI;
if (paddingMode == 1) { // in same paddingMode cudnn doesn't support asymmetric left/right top/bottopm paddings
auto ret = checkConv2dCUDNNPadAsymmetric(input, gradI, iH, iW, oH, oW, kH, kW, sH, sW, pH, pW, dH, dW, isNCHW);
tmpInput = std::move(std::get<0>(ret));
tmpGradI = std::move(std::get<1>(ret));
if (tmpInput) newInput = tmpInput.get();
if (tmpGradI) newGradI = tmpGradI.get();
}
conv2dBpCUDNN(block.launchContext(), newInput, newWeights, gradO, newGradI, newGradW, gradB, kH, kW, sH, sW, pH, pW,
dH, dW, paddingMode, isNCHW, wFormat);
if (0 == wFormat) {
// Create named vectors as lvalues
std::vector<LongType> nchwPermute = {2, 3, 1, 0};
std::vector<LongType> nhwcPermute = {1, 2, 3, 0};
// Use the appropriate one in the permutei call
newGradW->permutei(
isNCHW ? nchwPermute : nhwcPermute,false,false); // (oC, iC, kH, kW --> kH, kW, iC, oC) or (oC, kH, kW, iC --> kH, kW, iC, oC) iC, oC)
gradW->assign(newGradW);
}
if (newInput != input) {
if (isNCHW) {
NDArray assign = (*newGradI)({0, 0, 0, 0, 0, gradI->sizeAt(2), 0, gradI->sizeAt(3)});
gradI->assign(&assign);
} else {
NDArray assign = (*newGradI)({0, 0, 0, gradI->sizeAt(1), 0, gradI->sizeAt(2), 0, 0});
gradI->assign(&assign);
}
}
return Status::OK;
}
PLATFORM_CHECK(conv2d_bp, ENGINE_CUDA) {
auto input = INPUT_VARIABLE(0); // [bS, iH, iW, iC] (NHWC) or [bS, iC, iH, iW] (NCHW)
auto weights = INPUT_VARIABLE(1); // [kH, kW, iC, oC] always
auto bias = block.width() > 3 ? INPUT_VARIABLE(2) : nullptr; // [oC]
auto gradO = block.width() > 3
? INPUT_VARIABLE(3)
: INPUT_VARIABLE(2); // [bS, oH, oW, oC] (NHWC) or [bS, oC, oH, oW] (NCHW), epsilon_next
const int paddingMode = INT_ARG(8); // 0-VALID, 1-SAME, 2-CAUSAL
const int isNCHW = block.getIArguments()->size() > 9 ? !INT_ARG(9) : 1; // INT_ARG(9): 0-NCHW, 1-NHWC
Requirements req("CUDNN CONV2d_BP OP");
req.expectNotEq(makeInfoVariable(paddingMode, "paddingMode"), 2) &&
req.expectTrue(makeInfoVariable(isNCHW, "isNCHW")) &&
req.expectIn(makeInfoVariable(input->dataType(), TYPE_MSG_INPUT0),
{HALF, FLOAT32, DOUBLE}) &&
req.expectIn(makeInfoVariable(weights->dataType(), TYPE_MSG_INPUT1),
{HALF, FLOAT32, DOUBLE});
if (bias) {
req.expectIn(makeInfoVariable(bias->dataType(), TYPE_MSG_INPUT_ "#bias"),
{HALF, FLOAT32, DOUBLE}) &&
req.expectIn(makeInfoVariable(gradO->dataType(), TYPE_MSG_INPUT3),
{HALF, FLOAT32, DOUBLE});
} else {
req.expectIn(makeInfoVariable(gradO->dataType(), TYPE_MSG_INPUT2),
{HALF, FLOAT32, DOUBLE});
}
req.logTheSuccess();
return req;
}
} // namespace platforms
} // namespace ops
} // namespace sd
libnd4j/include/ops/declarable/platform/cudnn/conv3d.cu
+543
View File
@@ -0,0 +1,543 @@
/* ******************************************************************************
*
*
* This program and the accompanying materials are made available under the
* terms of the Apache License, Version 2.0 which is available at
* https://www.apache.org/licenses/LICENSE-2.0.
*
* See the NOTICE file distributed with this work for additional
* information regarding copyright ownership.
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
* WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
* License for the specific language governing permissions and limitations
* under the License.
*
* SPDX-License-Identifier: Apache-2.0
******************************************************************************/
//
// @author raver119@gmail.com
// @author Yurii Shyrma (iuriish@yahoo.com)
//
#include <ops/declarable/helpers/convolutions.h>
#include "cudnnUtils.h"
namespace sd {
namespace ops {
namespace platforms {
//////////////////////////////////////////////////////////////////////////
static void conv3dCUDNN(const LaunchContext* context, NDArray* input, NDArray* weights, NDArray* bias,
NDArray* output, const LongType kD, const LongType kH, const LongType kW, const LongType sD, const LongType sH,
const LongType sW, const LongType pD, const LongType pH, const LongType pW, const LongType dD, const LongType dH,
const LongType dW, const int paddingMode, const bool isNCDHW, const int wFormat) {
// cudnn support only one format for weights {oC,iC,kD,kH,kW}
const int numDims = 5;
LongType bS, iC, iD, iH, iW, oC, oD, oH,
oW; // batch size, input channels, input depth/height/width, output channels, output depth/height/width;
LongType indIOioC, indIOioD, indWoC, indWiC, indWkD; // corresponding indexes
ConvolutionUtils::getSizesAndIndexesConv3d(isNCDHW, wFormat, *input, *output, bS, iC, iD, iH, iW, oC, oD, oH, oW,
indIOioC, indIOioD, indWiC, indWoC, indWkD);
auto handle = reinterpret_cast<cudnnHandle_t*>(context->getCuDnnHandle());
CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnSetStream), cudnnSetStream(*handle, *context->getCudaStream()));
const std::vector<int> pads = {static_cast<int>(pD), static_cast<int>(pH), static_cast<int>(pW)};
const std::vector<int> filtStrides = {static_cast<int>(sD), static_cast<int>(sH), static_cast<int>(sW)};
const std::vector<int> dilations = {static_cast<int>(dD), static_cast<int>(dH), static_cast<int>(dW)};
const std::vector<int> xShape = {static_cast<int>(bS), static_cast<int>(iC), static_cast<int>(iD), static_cast<int>(iH), static_cast<int>(iW)};
const std::vector<int> zShape = {static_cast<int>(bS), static_cast<int>(oC), static_cast<int>(oD), static_cast<int>(oH), static_cast<int>(oW)};
const std::vector<int> wShape = {static_cast<int>(oC), static_cast<int>(iC), static_cast<int>(kD), static_cast<int>(kH), static_cast<int>(kW)};
const std::vector<int> bShape = {1, static_cast<int>(oC), 1, 1, 1};
const std::vector<int> xStrides = {static_cast<int>(input->strideAt(0)), static_cast<int>(input->strideAt(1)), static_cast<int>(input->strideAt(2)),
static_cast<int>(input->strideAt(3)), static_cast<int>(input->strideAt(4))};
const std::vector<int> zStrides = {static_cast<int>(output->strideAt(0)), static_cast<int>(output->strideAt(1)), static_cast<int>(output->strideAt(2)),
static_cast<int>(output->strideAt(3)), static_cast<int>(output->strideAt(4))};
cudnnTensorFormat_t format = isNCDHW ? CUDNN_TENSOR_NCHW : CUDNN_TENSOR_NHWC;
PointersManager manager(context, __func__);
// input descriptor
CudnnTensor x;
x.set(cudnnDataType(input->dataType()), numDims, xShape.data(), xStrides.data());
// weights descriptor
FilterDesc w;
w.set(cudnnDataType(weights->dataType()), CUDNN_TENSOR_NCHW, numDims, wShape.data());
// output descriptor
CudnnTensor z;
z.set(cudnnDataType(output->dataType()), numDims, zShape.data(), zStrides.data());
// description of convolution
ConvolutionDesc conv;
conv.set(numDims - 2, pads.data(), filtStrides.data(), dilations.data(), CUDNN_CROSS_CORRELATION,
cudnnDataType(output->dataType()));
// algorithm description
cudnnConvolutionFwdAlgo_t algo;
cudnnConvolutionFwdAlgoPerf_t algoPerf;
int count = 0;
CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnFindConvolutionForwardAlgorithm),
cudnnFindConvolutionForwardAlgorithm(*handle, x, w, conv, z, 1, &count, &algoPerf));
if (count == 0)
throw cuda_exception::build("conv3dCUDNN: cudnnGetConvolutionForwardAlgorithm failed as the count is 0", 0);
algo = algoPerf.algo;
// allocate auxiliary device memory, abbreviation ws means workspace
size_t wsSize;
CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnGetConvolutionForwardWorkspaceSize),
cudnnGetConvolutionForwardWorkspaceSize(*handle, x, w, conv, z, algo, &wsSize));
void* wsData = manager.allocateDevMem(wsSize);
// provide scaling parameters
const float alpha32(1), beta32(0);
const double alpha64(1), beta64(0);
const void* alpha =
output->sizeOfT() <= 4 ? reinterpret_cast<const void*>(&alpha32) : reinterpret_cast<const void*>(&alpha64);
const void* beta =
output->sizeOfT() <= 4 ? reinterpret_cast<const void*>(&beta32) : reinterpret_cast<const void*>(&beta64);
NDArray::prepareSpecialUse({output}, {input, weights, bias});
// run calculation
CHECK_CUDNN_FAILURE_MSG(
STRINGIZE(cudnnConvolutionForward),
cudnnConvolutionForward(*handle, alpha, x, input->specialBuffer(), w, weights->specialBuffer(), conv, algo,
wsData, wsSize, beta, z, output->specialBuffer()));
// add bias if it is present
if (bias != nullptr) {
CudnnTensor b;
b.setEx(/*format*/ CUDNN_TENSOR_NCHW, cudnnDataType(bias->dataType()), numDims, bShape.data());
CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnAddTensor), cudnnAddTensor(*handle, alpha, b, bias->specialBuffer(), alpha,
z, output->specialBuffer()));
}
NDArray::registerSpecialUse({output}, {input, weights, bias});
}
//////////////////////////////////////////////////////////////////////////
static void conv3dBpCUDNN(const LaunchContext* context, NDArray* input, NDArray* weights,
NDArray* gradO, NDArray* gradI, NDArray* gradW, NDArray* gradB, const int kD,
const LongType kH, const LongType kW, const LongType sD, const LongType sH, const LongType sW, const LongType pD,
const LongType pH, const LongType pW, const LongType dD, const LongType dH, const LongType dW, const int paddingMode,
const bool isNCDHW, const int wFormat) {
// cudnn supports only two formats {oC,iC,kD,kH,kW} and {oC,kD,kH,kW,iC} for weights/gradW
const int numDims = 5;
LongType bS, iC, iD, iH, iW, oC, oD, oH,
oW; // batch size, input channels, input depth/height/width, output channels, output depth/height/width;
LongType indIOioC, indIOioD, indWoC, indWiC, indWkD; // corresponding indexes
ConvolutionUtils::getSizesAndIndexesConv3d(isNCDHW, wFormat, *input, *gradO, bS, iC, iD, iH, iW, oC, oD, oH, oW,
indIOioC, indIOioD, indWiC, indWoC, indWkD);
auto handle = reinterpret_cast<cudnnHandle_t*>(context->getCuDnnHandle());
CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnSetStream), cudnnSetStream(*handle, *context->getCudaStream()));
const std::vector<int> pads = {static_cast<int>(pD), static_cast<int>(pH), static_cast<int>(pW)};
const std::vector<int> filtStrides = {static_cast<int>(sD), static_cast<int>(sH), static_cast<int>(sW)};
const std::vector<int> dilations = {static_cast<int>(dD), static_cast<int>(dH), static_cast<int>(dW)};
const std::vector<int> xShape = {static_cast<int>(bS), static_cast<int>(iC), static_cast<int>(iD), static_cast<int>(iH), static_cast<int>(iW)};
const std::vector<int> dzShape = {static_cast<int>(bS), static_cast<int>(oC), static_cast<int>(oD), static_cast<int>(oH), static_cast<int>(oW)};
const std::vector<int> wShape = {static_cast<int>(oC), static_cast<int>(iC), static_cast<int>(kD), static_cast<int>(kH), static_cast<int>(kW)};
const std::vector<int> dbShape = {1, static_cast<int>(isNCDHW ? oC : 1), 1, 1, (int)(isNCDHW ? 1 : oC)};
const std::vector<int> xStrides = {static_cast<int>(input->strideAt(0)), static_cast<int>(input->strideAt(1)), static_cast<int>(input->strideAt(2)),
static_cast<int>(input->strideAt(3)), static_cast<int>(input->strideAt(4))};
const std::vector<int> dxStrides = {static_cast<int>(gradI->strideAt(0)), static_cast<int>(gradI->strideAt(1)), static_cast<int>(gradI->strideAt(2)),
static_cast<int>(gradI->strideAt(3)), static_cast<int>(gradI->strideAt(4))};
const std::vector<int> dzStrides = {static_cast<int>(gradO->strideAt(0)), static_cast<int>(gradO->strideAt(1)), static_cast<int>(gradO->strideAt(2)),
static_cast<int>(gradO->strideAt(3)), static_cast<int>(gradO->strideAt(4))};
cudnnTensorFormat_t format = isNCDHW ? CUDNN_TENSOR_NCHW : CUDNN_TENSOR_NHWC;
cudnnTensorFormat_t formatW = 0 == wFormat ? format : (1 == wFormat ? CUDNN_TENSOR_NCHW : CUDNN_TENSOR_NHWC);
PointersManager manager(context, __func__);
// input descriptor, gradO descriptor, gradI descriptor
CudnnTensor x, dz, dx;
x.set(cudnnDataType(input->dataType()), numDims, xShape.data(), xStrides.data());
dz.set(cudnnDataType(gradO->dataType()), numDims, dzShape.data(), dzStrides.data());
dx.set(cudnnDataType(gradI->dataType()), numDims, xShape.data(), dxStrides.data());
// gradW descriptor
FilterDesc dw;
dw.set(cudnnDataType(gradW->dataType()), formatW, numDims, wShape.data());
// description of convolution
ConvolutionDesc conv;
conv.set(numDims - 2, pads.data(), filtStrides.data(), dilations.data(), CUDNN_CROSS_CORRELATION,
cudnnDataType(gradO->dataType()));
// gradW algorithm description
cudnnConvolutionBwdFilterAlgo_t algoGradW;
cudnnConvolutionBwdFilterAlgoPerf_t algoGradWPerf;
int count = 0;
CHECK_CUDNN_FAILURE_MSG(
STRINGIZE(cudnnFindConvolutionBackwardFilterAlgorithm),
cudnnFindConvolutionBackwardFilterAlgorithm(*handle, x, dz, conv, dw, 1, &count, &algoGradWPerf));
if (count == 0)
throw cuda_exception::build(
"conv3dBpCUDNN: cudnnGetConvolutionBackwardFilterAlgorithm failed as the count is 0", 0);
algoGradW = algoGradWPerf.algo;
// gradI algorithm description
cudnnConvolutionBwdDataAlgo_t algoGradI;
cudnnConvolutionBwdDataAlgoPerf_t algoGradIPerf;
// CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnGetConvolutionBackwardDataAlgorithm),
// cudnnGetConvolutionBackwardDataAlgorithm( *handle, dw, dz, conv, x, CUDNN_CONVOLUTION_BWD_DATA_PREFER_FASTEST, 0,
// &algoGradI));
CHECK_CUDNN_FAILURE_MSG(
STRINGIZE(cudnnFindConvolutionBackwardDataAlgorithm),
cudnnFindConvolutionBackwardDataAlgorithm(*handle, dw, dz, conv, x, 1, &count, &algoGradIPerf));
if (count == 0)
throw cuda_exception::build("conv3dBpCUDNN: cudnnGetConvolutionBackwardDataAlgorithm failed as the count is 0",
0);
algoGradI = algoGradIPerf.algo;
// allocate auxiliary device memory for gradW calculation, abbreviation ws means workspace
size_t wsGradWSize;
CHECK_CUDNN_FAILURE_MSG(
STRINGIZE(cudnnGetConvolutionBackwardFilterWorkspaceSize),
cudnnGetConvolutionBackwardFilterWorkspaceSize(*handle, x, dz, conv, dw, algoGradW, &wsGradWSize));
void* wsGradWData = manager.allocateDevMem(wsGradWSize);
// allocate auxiliary device memory for gradI calculation, abbreviation ws means workspace
size_t wsGradISize;
CHECK_CUDNN_FAILURE_MSG(
STRINGIZE(cudnnGetConvolutionBackwardDataWorkspaceSize),
cudnnGetConvolutionBackwardDataWorkspaceSize(*handle, dw, dz, conv, dx, algoGradI, &wsGradISize));
void* wsGradIData = manager.allocateDevMem(wsGradISize);
// provide scaling parameters
const float alpha32(1), beta32(0);
const double alpha64(1), beta64(0);
const void* alpha =
gradO->sizeOfT() <= 4 ? reinterpret_cast<const void*>(&alpha32) : reinterpret_cast<const void*>(&alpha64);
const void* beta =
gradO->sizeOfT() <= 4 ? reinterpret_cast<const void*>(&beta32) : reinterpret_cast<const void*>(&beta64);
NDArray::prepareSpecialUse({gradI, gradW, gradB}, {input, weights, gradO});
// run calculation for gradB (if not nullptr)
if (gradB != nullptr) {
CudnnTensor db;
db.setEx(format, cudnnDataType(gradB->dataType()), numDims, dbShape.data());
CHECK_CUDNN_FAILURE_MSG(
STRINGIZE(cudnnConvolutionBackwardBias),
cudnnConvolutionBackwardBias(*handle, alpha, dz, gradO->specialBuffer(), beta, db, gradB->specialBuffer()));
}
// run calculation for gradW
CHECK_CUDNN_FAILURE_MSG(
STRINGIZE(cudnnConvolutionBackwardFilter),
cudnnConvolutionBackwardFilter(*handle, alpha, x, input->specialBuffer(), dz, gradO->specialBuffer(), conv,
algoGradW, wsGradWData, wsGradWSize, beta, dw, gradW->specialBuffer()));
// run calculation for gradI
CHECK_CUDNN_FAILURE_MSG(
STRINGIZE(cudnnConvolutionBackwardData),
cudnnConvolutionBackwardData(*handle, alpha, dw, weights->specialBuffer(), dz, gradO->specialBuffer(), conv,
algoGradI, wsGradIData, wsGradISize, beta, dx, gradI->specialBuffer()));
NDArray::registerSpecialUse({gradI, gradW, gradB}, {input, weights, gradO});
}
//////////////////////////////////////////////////////////////////////////
PLATFORM_IMPL(conv3dnew, ENGINE_CUDA) {
auto input = INPUT_VARIABLE(0); // [bS, iD, iH, iW, iC] (NDHWC) or [bS, iC, iD, iH, iW] (NCDHW)
auto weights = INPUT_VARIABLE(1); // [kD, kH, kW, iC, oC], [oC, iC, kD, kH, kW], [oC, kD, kH, kW, iC]
auto bias = block.width() > 2 ? INPUT_VARIABLE(2) : nullptr; // [oC]
auto output = OUTPUT_VARIABLE(0); // [bS, oD, oH, oW, oC] (NDHWC) or [bS, oC, oD, oH, oW] (NCDHW)
REQUIRE_TRUE(input->rankOf() == 5, 0, "CONV3D CUDNN OP: rank of input array must be equal to 5, but got %i instead !",
input->rankOf());
REQUIRE_TRUE(weights->rankOf() == 5, 0,
"CONV3D CUDNN OP: rank of weights array must be equal to 5, but got %i instead !", weights->rankOf());
LongType kD = INT_ARG(0) > 0 ? INT_ARG(0) : static_cast<LongType>(weights->sizeAt(0)); // filter(kernel) depth
LongType kH = INT_ARG(1) > 0 ? INT_ARG(1) : static_cast<LongType>(weights->sizeAt(1)); // filter(kernel) height
LongType kW = INT_ARG(2) > 0 ? INT_ARG(2) : static_cast<LongType>(weights->sizeAt(2)); // filter(kernel) width
LongType sD = INT_ARG(3); // strides depth
LongType sH = INT_ARG(4); // strides height
LongType sW = INT_ARG(5); // strides width
LongType pD = INT_ARG(6); // paddings depth
LongType pH = INT_ARG(7); // paddings height
LongType pW = INT_ARG(8); // paddings width
LongType dD = INT_ARG(9); // dilations depth
LongType dH = INT_ARG(10); // dilations height
LongType dW = INT_ARG(11); // dilations width
int paddingMode = INT_ARG(12); // 0-SAME, 1-VALID
int isNCDHW = block.getIArguments()->size() > 13 ? !INT_ARG(13) : 1; // INT_ARG(13): 1-NDHWC, 0-NCDHW
int wFormat = block.getIArguments()->size() > 14
? INT_ARG(14)
: 0; // 0-[kD, kH, kW, iC, oC], 1-[oC, iC, kD, kH, kW], 2-[oC, kD, kH, kW, iC]
REQUIRE_TRUE(paddingMode < 2, 0,
"CONV3D CUDNN OP: causal padding mode (paddingMode = 2) is not allowed for this operation !");
LongType bS, iC, iD, iH, iW, oC, oD, oH,
oW; // batch size, input channels, input depth/height/width, output channels, output depth/height/width;
LongType indIOioC, indIOioD, indWoC, indWiC, indWkD; // corresponding indexes
ConvolutionUtils::getSizesAndIndexesConv3d(isNCDHW, wFormat, *input, *output, bS, iC, iD, iH, iW, oC, oD, oH, oW,
indIOioC, indIOioD, indWiC, indWoC, indWkD);
ConvolutionUtils::calcPadding3D(pD, pH, pW, oD, oH, oW, iD, iH, iW, kD, kH, kW, sD, sH, sW, dD, dH, dW, paddingMode);
std::vector<LongType> expectedWeightsShape = ConvolutionUtils::expectWeightsShape(wFormat, kD, kH, kW, iC, oC);
REQUIRE_TRUE(weights->isSameShape(expectedWeightsShape), 0,
"CONV3D CUDNN OP: wrong shape of weights array, expected is %s, but got %s instead !",
ShapeUtils::shapeAsString(expectedWeightsShape).c_str(), ShapeUtils::shapeAsString(weights).c_str());
if (bias)
REQUIRE_TRUE(
bias->rankOf() <= 2 && oC == bias->lengthOf(), 0,
"CONV3D CUDNN OP: wrong shape of array with biases, expected rank, length: <=2, %i, but got %i, %i instead !",
oC, bias->rankOf(), bias->lengthOf());
std::unique_ptr<NDArray> tmpWeight = {}, tmpInput = {};
NDArray* newWeights = weights; // cudnn support only one format {oC,iC,kD,kH,kW}
if (1 != wFormat) {
// Create the tmpWeight object - this syntax is already valid
std::vector<LongType> weightShape = {oC, iC, kD, kH, kW};
// Use the vector for the NDArray constructor
tmpWeight.reset(new NDArray(weights->ordering(), weightShape, weights->dataType(), weights->getContext()));
newWeights = tmpWeight.get();
// Create named vectors as lvalues for the permute call
std::vector<LongType> format0Permute = {4, 3, 0, 1, 2};
std::vector<LongType> format1Permute = {0, 4, 1, 2, 3};
NDArray assign = weights->permute(
0 == wFormat ? format0Permute : format1Permute,
true, // copyToNewBuff
true);
// Use the appropriate one in the permute call
newWeights->assign(&assign); // resetStrides
}
if (paddingMode == 1) { // in same paddingMode cudnn doesn't support asymmetric left/right top/bottopm paddings
auto ret = checkConv3dCUDNNPadAsymmetric(input, nullptr, iD, iH, iW, oD, oH, oW, kD, kH, kW, sD, sH, sW, pD, pH, pW,
dD, dH, dW, isNCDHW);
tmpInput = std::move(std::get<0>(ret)); // prolong life
if (tmpInput) input = tmpInput.get();
}
conv3dCUDNN(block.launchContext(), input, newWeights, bias, output, kD, kH, kW, sD, sH, sW, pD, pH, pW, dD, dH, dW,
paddingMode, isNCDHW, wFormat);
return Status::OK;
}
//////////////////////////////////////////////////////////////////////////
PLATFORM_CHECK(conv3dnew, ENGINE_CUDA) {
auto input = INPUT_VARIABLE(0); // [bS, iD, iH, iW, iC] (NDHWC) or [bS, iC, iD, iH, iW] (NCDHW)
auto weights = INPUT_VARIABLE(1); // [kD, kH, kW, iC, oC], [oC, iC, kD, kH, kW], [oC, kD, kH, kW, iC]
auto bias = block.width() > 2 ? INPUT_VARIABLE(2) : nullptr; // [oC]
int paddingMode = INT_ARG(12); // 0-SAME, 1-VALID
Requirements req("CUDNN CONV3d OP");
req.expectNotEq(makeInfoVariable(paddingMode, "paddingMode"), 2) &&
req.expectIn(makeInfoVariable(input->dataType(), TYPE_MSG_INPUT0),
{HALF, FLOAT32, DOUBLE}) &&
req.expectIn(makeInfoVariable(weights->dataType(), TYPE_MSG_INPUT1),
{HALF, FLOAT32, DOUBLE});
if (bias) {
req.expectIn(makeInfoVariable(bias->dataType(), TYPE_MSG_INPUT_ "#bias"),
{HALF, FLOAT32, DOUBLE});
}
req.logTheSuccess();
return req;
}
//////////////////////////////////////////////////////////////////////////
PLATFORM_IMPL(conv3dnew_bp, ENGINE_CUDA) {
auto input = INPUT_VARIABLE(0); // [bS, iD, iH, iW, iC] (NDHWC) or [bS, iC, iD, iH, iW] (NCDHW)
auto weights = INPUT_VARIABLE(1); // [kD, kH, kW, iC, oC], [oC, iC, kD, kH, kW], [oC, kD, kH, kW, iC]
auto bias = block.width() > 3 ? INPUT_VARIABLE(2) : nullptr; // [oC]
auto gradO = block.width() > 3
? INPUT_VARIABLE(3)
: INPUT_VARIABLE(2); // [bS, oD, oH, oW, oC] (NDHWC) or [bS, oC, oD, oH, oW] (NCDHW), epsilon_next
auto gradI = OUTPUT_VARIABLE(0); // [bS, iD, iH, iW, iC] (NDHWC) or [bS, iC, iD, iH, iW] (NCDHW), epsilon
auto gradW = OUTPUT_VARIABLE(1); // [kD, kH, kW, iC, oC], [oC, iC, kD, kH, kW], [oC, kD, kH, kW, iC]
auto gradB = block.width() > 3 ? OUTPUT_VARIABLE(2) : nullptr; // [oC]
REQUIRE_TRUE(input->rankOf() == 5, 0,
"CONV3D_BP CUDNN OP: rank of input array must be equal to 5, but got %i instead !", input->rankOf());
REQUIRE_TRUE(weights->rankOf() == 5, 0,
"CONV3D_BP CUDNN OP: rank of weights array must be equal to 5, but got %i instead !", weights->rankOf());
REQUIRE_TRUE(
gradO->rankOf() == 5, 0,
"CONV3D_BP CUDNN OP: rank of output gradients (next epsilon) array must be equal to 5, but got %i instead !",
gradO->rankOf());
LongType kD = INT_ARG(0) > 0 ? INT_ARG(0) : static_cast<LongType>(weights->sizeAt(0)); // filter(kernel) depth
LongType kH = INT_ARG(1) > 0 ? INT_ARG(1) : static_cast<LongType>(weights->sizeAt(1)); // filter(kernel) height
LongType kW = INT_ARG(2) > 0 ? INT_ARG(2) : static_cast<LongType>(weights->sizeAt(2)); // filter(kernel) width
LongType sD = INT_ARG(3); // strides depth
LongType sH = INT_ARG(4); // strides height
LongType sW = INT_ARG(5); // strides width
LongType pD = INT_ARG(6); // paddings depth
LongType pH = INT_ARG(7); // paddings height
LongType pW = INT_ARG(8); // paddings width
LongType dD = INT_ARG(9); // dilations depth
LongType dH = INT_ARG(10); // dilations height
LongType dW = INT_ARG(11); // dilations width
int paddingMode = INT_ARG(12); // 1-SAME, 0-VALID
int isNCDHW = block.getIArguments()->size() > 13 ? !INT_ARG(13) : 1; // INT_ARG(13): 1-NDHWC, 0-NCDHW
int wFormat = block.getIArguments()->size() > 14
? INT_ARG(14)
: 0; // 0-[kD, kH, kW, iC, oC], 1-[oC, iC, kD, kH, kW], 2-[oC, kD, kH, kW, iC]
LongType bS, iC, iD, iH, iW, oC, oD, oH,
oW; // batch size, input channels, input depth/height/width, output channels, output depth/height/width;
LongType indIOioC, indIOioD, indWoC, indWiC, indWkD; // corresponding indexes
ConvolutionUtils::getSizesAndIndexesConv3d(isNCDHW, wFormat, *input, *gradO, bS, iC, iD, iH, iW, oC, oD, oH, oW,
indIOioC, indIOioD, indWiC, indWoC, indWkD);
LongType trueoD, trueoH, trueoW; // true output depth/height/width
ConvolutionUtils::calcOutSizePool3D(trueoD, trueoH, trueoW, kD, kH, kW, sD, sH, sW, pD, pH, pW, dD, dH, dW, iD, iH,
iW, paddingMode);
REQUIRE_TRUE(paddingMode < 2, 0,
"CONV3D_BP CUDNN OP: causal padding mode (paddingMode = 2) is not allowed for this operation !");
std::vector<LongType> expectedGradOShape = ShapeUtils::composeShapeUsingDimsAndIdx(
{bS, oC, trueoD, trueoH, trueoW, 0, indIOioC, indIOioD, indIOioD + 1, indIOioD + 2});
std::vector<LongType> expectedWeightsShape = ConvolutionUtils::expectWeightsShape(wFormat, kD, kH, kW, iC, oC);
REQUIRE_TRUE(
gradO->isSameShape(expectedGradOShape), 0,
"CONV3D_BP CUDNN OP: wrong shape of output gradients (next epsilon) array, expected is %s, but got %s instead !",
ShapeUtils::shapeAsString(expectedGradOShape).c_str(), ShapeUtils::shapeAsString(gradO).c_str());
REQUIRE_TRUE(gradW->isSameShape(expectedWeightsShape), 0,
"CONV3D_BP CUDNN OP: wrong shape of weights array, expected is %s, but got %s instead !",
ShapeUtils::shapeAsString(expectedWeightsShape).c_str(), ShapeUtils::shapeAsString(weights).c_str());
if (bias)
REQUIRE_TRUE(bias->rankOf() <= 2 && oC == bias->lengthOf(), 0,
"CONV3D_BP CUDNN OP: wrong shape of array with biases, expected rank, length: <=2, %i, but got %i, %i "
"instead !",
oC, bias->rankOf(), bias->lengthOf());
ConvolutionUtils::calcPadding3D(pD, pH, pW, oD, oH, oW, iD, iH, iW, kD, kH, kW, sD, sH, sW, dD, dH, dW, paddingMode);
std::unique_ptr<NDArray> tmpGradI = {}, tmpInput = {}, tmpWeights = {}, tmpGradW = {};
NDArray *newWeights = weights,
*newGradW = gradW; // cudnn support only two formats {oC,iC,kD,kH,kW} and {oC,kD,kH,kW,iC}
if (0 == wFormat) {
// Create named vectors as lvalues for the NDArray constructor
std::vector<LongType> ncdhwGradShape = {oC, iC, kD, kH, kW};
std::vector<LongType> ndhwcGradShape = {oC, kD, kH, kW, iC};
// Use the appropriate one for the gradW reset
tmpGradW.reset(new NDArray(
gradW->ordering(),
isNCDHW ? ncdhwGradShape : ndhwcGradShape,
gradW->dataType(), gradW->getContext()));
// Create named vectors as lvalues for the NDArray constructor
std::vector<LongType> ncdhwShape = {oC, iC, kD, kH, kW};
std::vector<LongType> ndhwcShape = {oC, kD, kH, kW, iC};
// Use the appropriate one for the weights reset
tmpWeights.reset(new NDArray(
weights->ordering(),
isNCDHW ? ncdhwShape : ndhwcShape,
weights->dataType(), weights->getContext()));
// Create named vectors as lvalues for the permute call
std::vector<LongType> ncdhwPermute = {4, 3, 0, 1, 2};
std::vector<LongType> ndhwcPermute = {4, 0, 1, 2, 3};
// Set the pointer variables
newGradW = tmpGradW.get();
newWeights = tmpWeights.get();
NDArray assign = weights->permute(
isNCDHW ? ncdhwPermute : ndhwcPermute,
true, // copyToNewBuff
true);
// Use the appropriate one in the permute call
newWeights->assign(&assign); // resetStrides (kD, kH, kW, iC, oC --> oC, iC, kD, kH, kW) or (kD, kH, kW,
// iC, oC --> oC, kD, kH, kW, iC)
}
NDArray* newInput = input;
NDArray* newGradI = gradI;
if (paddingMode == 1) { // in same paddingMode cudnn doesn't support asymmetric left/right top/bottopm paddings
auto ret = checkConv3dCUDNNPadAsymmetric(input, gradI, iD, iH, iW, oD, oH, oW, kD, kH, kW, sD, sH, sW, pD, pH, pW,
dD, dH, dW, isNCDHW);
tmpInput = std::move(std::get<0>(ret));
tmpGradI = std::move(std::get<1>(ret));
if (tmpInput) newInput = tmpInput.get();
if (tmpGradI) newGradI = tmpGradI.get();
}
conv3dBpCUDNN(block.launchContext(), newInput, newWeights, gradO, newGradI, newGradW, gradB, kD, kH, kW, sD, sH, sW,
pD, pH, pW, dD, dH, dW, paddingMode, isNCDHW, wFormat);
if (0 == wFormat) {
// Create named vectors as lvalues for the permutei call
std::vector<LongType> ncdhwPermutei = {2, 3, 4, 1, 0};
std::vector<LongType> ndhwcPermutei = {1, 2, 3, 4, 0};
// Use the appropriate one in the permutei call
newGradW->permutei(
isNCDHW ? ncdhwPermutei : ndhwcPermutei,false,false); // (oC, iC, kD, kH, kW --> kD, kH, kW, iC, oC) or (oC, kD, kH, kW, iC --> kD, kH, kW, iC, oC) gradW->assign(newGradW);
}
if (newInput != input) {
if (isNCDHW) {
NDArray assign = (*newGradI)({0, 0, 0, 0, 0, gradI->sizeAt(2), 0, gradI->sizeAt(3), 0, gradI->sizeAt(4)});
gradI->assign(&assign);
} else {
NDArray assign = (*newGradI)({0, 0, 0, gradI->sizeAt(1), 0, gradI->sizeAt(2), 0, gradI->sizeAt(3), 0, 0});
gradI->assign(&assign);
}
}
return Status::OK;
}
PLATFORM_CHECK(conv3dnew_bp, ENGINE_CUDA) {
auto input = INPUT_VARIABLE(0); // [bS, iD, iH, iW, iC] (NDHWC) or [bS, iC, iD, iH, iW] (NCDHW)
auto weights = INPUT_VARIABLE(1); // [kD, kH, kW, iC, oC], [oC, iC, kD, kH, kW], [oC, kD, kH, kW, iC]
auto bias = block.width() > 3 ? INPUT_VARIABLE(2) : nullptr; // [oC]
auto gradO = block.width() > 3
? INPUT_VARIABLE(3)
: INPUT_VARIABLE(2); // [bS, oD, oH, oW, oC] (NDHWC) or [bS, oC, oD, oH, oW] (NCDHW), epsilon_next
LongType paddingMode = INT_ARG(12); // 1-SAME, 0-VALID
int isNCDHW = block.getIArguments()->size() > 13 ? !INT_ARG(13) : 1; // INT_ARG(13): 1-NDHWC, 0-NCDHW
Requirements req("CUDNN CONV3d_BP OP");
req.expectNotEq(makeInfoVariable(paddingMode, "paddingMode"), 2) &&
req.expectTrue(makeInfoVariable(isNCDHW, "isNCDHW")) &&
req.expectIn(makeInfoVariable(input->dataType(), TYPE_MSG_INPUT0),
{HALF, FLOAT32, DOUBLE}) &&
req.expectIn(makeInfoVariable(weights->dataType(), TYPE_MSG_INPUT1),
{HALF, FLOAT32, DOUBLE});
if (bias) {
req.expectIn(makeInfoVariable(bias->dataType(), TYPE_MSG_INPUT_ "#bias"),
{HALF, FLOAT32, DOUBLE}) &&
req.expectIn(makeInfoVariable(gradO->dataType(), TYPE_MSG_INPUT3),
{HALF, FLOAT32, DOUBLE});
} else {
req.expectIn(makeInfoVariable(gradO->dataType(), TYPE_MSG_INPUT2),
{HALF, FLOAT32, DOUBLE});
}
req.logTheSuccess();
return req;
}
} // namespace platforms
} // namespace ops
} // namespace sd
libnd4j/include/ops/declarable/platform/cudnn/ctcloss.cu
+189
View File
@@ -0,0 +1,189 @@
/*******************************************************************************
*
* Copyright (c) 2021 Konduit K.K.
*
* This program and the accompanying materials are made available under the
* terms of the Apache License, Version 2.0 which is available at
* https://www.apache.org/licenses/LICENSE-2.0.
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
* WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
* License for the specific language governing permissions and limitations
* under the License.
*
* SPDX-License-Identifier: Apache-2.0
******************************************************************************/
//
// @author AbdelRauf
//
#include <array/NDArrayFactory.h>
#include <vector>
#include "cudnnUtils.h"
namespace sd {
namespace ops {
namespace platforms {
std::vector<int> getConcatTargets(NDArray&targetLabels, NDArray&targetLabelLengths) {
// concatenate target labels
const int32_t *tlabels = bufferInHost<int32_t>(targetLabels);
const int32_t *tlens = bufferInHost<int32_t>(targetLabelLengths);
int32_t nextOffset = targetLabels.strideAt(0);
int32_t elStride = targetLabels.strideAt(1);
int32_t batchCount = targetLabelLengths.lengthOf();
std::vector<int> labels;
labels.resize(targetLabels.lengthOf());
int j = 0;
for (int i = 0; i < batchCount; i++) {
int count = tlens[i];
for (int k = 0; k < count; k++) {
labels[j] = tlabels[k * elStride];
j++;
}
tlabels += nextOffset;
}
return labels;
}
void cudnnCtcLoss(const LaunchContext &context, NDArray&probs, const int32_t *targetLabelsPtr,
NDArray&probInputLengthes, NDArray&targetLabelLengths, NDArray &ctcLosses,
NDArray &grads) {
const int dims[] = {(int)probs.sizeAt(0), (int)probs.sizeAt(1), (int)probs.sizeAt(2)};
const int strides[] = {(int)probs.strideAt(0), (int)probs.strideAt(1), (int)probs.strideAt(2)};
auto handle = reinterpret_cast<cudnnHandle_t *>(context.getCuDnnHandle());
CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnSetStream), cudnnSetStream(*handle, *context.getCudaStream()));
CTCLossDesc ctcLossDesc;
CudnnTensor probsDesc, gradsDesc(nullptr);
bool calcGrads = !grads.isEmpty();
auto cudnnType = cudnnDataType(probs.dataType());
ctcLossDesc.set(cudnnType, CUDNN_LOSS_NORMALIZATION_SOFTMAX, CUDNN_PROPAGATE_NAN);
probsDesc.set(cudnnType, probs.rankOf(), dims, strides);
if (calcGrads) {
gradsDesc.create();
const int gradStrides[] = {(int)grads.strideAt(0), (int)grads.strideAt(1), (int)grads.strideAt(2)};
gradsDesc.set(cudnnDataType(grads.dataType()), grads.rankOf(), dims, gradStrides);
}
size_t tempWorkSpaceSize = 0;
CHECK_CUDNN_FAILURE_MSG(
STRINGIZE(cudnnGetCTCLossWorkspaceSize),
cudnnGetCTCLossWorkspaceSize(*handle, probsDesc, gradsDesc, targetLabelsPtr,
bufferInHost<int32_t>(targetLabelLengths), bufferInHost<int32_t>(probInputLengthes),
CUDNN_CTC_LOSS_ALGO_DETERMINISTIC, ctcLossDesc, &tempWorkSpaceSize));
PointersManager manager(&context, __func__);
// Allocate temp tempWorkspace buffer
void *tempWorkSpace = manager.allocateDevMem(tempWorkSpaceSize);
NDArray::prepareSpecialUse({&ctcLosses, &grads}, {&probs});
CHECK_CUDNN_FAILURE_MSG(
STRINGIZE(cudnnCTCLoss),
cudnnCTCLoss(*handle, probsDesc, probs.specialBuffer(), targetLabelsPtr,
bufferInHost<int32_t>(targetLabelLengths), bufferInHost<int32_t>(probInputLengthes),
ctcLosses.specialBuffer(), gradsDesc, calcGrads ? grads.specialBuffer() : nullptr,
CUDNN_CTC_LOSS_ALGO_DETERMINISTIC, ctcLossDesc, tempWorkSpace, tempWorkSpaceSize));
NDArray::registerSpecialUse({&ctcLosses, &grads}, {&probs});
return;
}
PLATFORM_IMPL(ctc_loss, ENGINE_CUDA) {
auto targetLabels = INPUT_VARIABLE(0);
auto logitInput = INPUT_VARIABLE(1);
auto targetLabelLengths = INPUT_VARIABLE(2);
auto logitInputLengths = INPUT_VARIABLE(3);
auto outputLosses = OUTPUT_VARIABLE(0);
auto context = block.launchContext();
// in Cudnn Batch is in the middle dimension
logitInput->permutei({1, 0, 2});
// in Cudnn targets are concantenated instead of batched as matrix
auto labels = getConcatTargets(*targetLabels, *targetLabelLengths);
const int32_t *ldata = labels.data();
auto emptyGrads = NDArrayFactory::empty<float>();
cudnnCtcLoss(*context, *logitInput, ldata, *logitInputLengths, *targetLabelLengths, *outputLosses, emptyGrads);
return Status::OK;
}
template <typename T>
bool checkLabelLength(NDArray&labelLengthArr) {
// check label lengths
auto lenBatch = labelLengthArr.lengthOf();
for (int i = 0; i < lenBatch; i++) {
// The labelLengths is greater than 256.
if (labelLengthArr.e<int32_t>(i) > 256) return false;
}
return true;
}
PLATFORM_CHECK(ctc_loss, ENGINE_CUDA) {
auto targetLabels = INPUT_VARIABLE(0);
auto logitInput = INPUT_VARIABLE(1);
auto targetLabelLengths = INPUT_VARIABLE(2);
auto logitInputLengths = INPUT_VARIABLE(3);
auto outputLosses = OUTPUT_VARIABLE(0);
int blankIndex = INT_ARG(0);
Requirements req("CUDNN CTC_LOSS OP");
req.expectEq(makeInfoVariable(blankIndex, "Blank Index"), 0) &&
req.expectEq(makeInfoVariable(logitInput->dataType(), TYPE_MSG_INPUT1), FLOAT32) &&
req.expectEq(makeInfoVariable(targetLabelLengths->dataType(), TYPE_MSG_INPUT2), INT32) &&
req.expectTrue(
makeInfoVariable(checkLabelLength<int32_t>(*targetLabelLengths), "target Label lengthes should be <= 256"),
NO_MSG);
req.logTheSuccess();
return req;
}
PLATFORM_IMPL(ctc_loss_grad, ENGINE_CUDA) {
auto targetLabels = INPUT_VARIABLE(0);
auto logitInput = INPUT_VARIABLE(1);
auto targetLabelLengths = INPUT_VARIABLE(2);
auto logitInputLengths = INPUT_VARIABLE(3);
auto outputGradients = OUTPUT_VARIABLE(0);
auto context = block.launchContext();
REQUIRE_TRUE(outputGradients->isSameShape(logitInput), 0,
"CtcLoss Gradient: wrong shape of output array, expected is %s but got %s instead !",
ShapeUtils::shapeAsString(logitInput).c_str(), ShapeUtils::shapeAsString(outputGradients).c_str());
// in Cudnn Batch is in the middle dimension
logitInput->permutei({1, 0, 2});
outputGradients->permutei({1, 0, 2});
// in Cudnn targets are concantenated instead of batched as matrix
auto labels = getConcatTargets(*targetLabels, *targetLabelLengths);
const int32_t *ldata = labels.data();
auto tempLosses = NDArrayFactory::create<float>('c', {logitInputLengths->sizeAt(0)});
cudnnCtcLoss(*context, *logitInput, ldata, *logitInputLengths, *targetLabelLengths, tempLosses, *outputGradients);
// restore grads shape from {T, BATCH, C} -> {BATCHS, T, C}
outputGradients->permutei({1, 0, 2});
return Status::OK;
}
PLATFORM_CHECK(ctc_loss_grad, ENGINE_CUDA) {
auto targetLabels = INPUT_VARIABLE(0);
auto logitInput = INPUT_VARIABLE(1);
auto targetLabelLengths = INPUT_VARIABLE(2);
auto logitInputLengths = INPUT_VARIABLE(3);
auto outputGrads = OUTPUT_VARIABLE(0);
int blankIndex = INT_ARG(0);
Requirements req("CUDNN CTC_LOSS_GRAD OP");
req.expectEq(makeInfoVariable(blankIndex, "Blank Index"), 0) &&
req.expectEq(makeInfoVariable(logitInput->dataType(), TYPE_MSG_INPUT1), FLOAT32) &&
req.expectEq(makeInfoVariable(targetLabelLengths->dataType(), TYPE_MSG_INPUT2), INT32) &&
req.expectTrue(
makeInfoVariable(checkLabelLength<int32_t>(*targetLabelLengths), "target Label lengthes should be <= 256"),
NO_MSG);
req.logTheSuccess();
return req;
}
} // namespace platforms
} // namespace ops
} // namespace sd
libnd4j/include/ops/declarable/platform/cudnn/cudnnUtils.cu
+397
View File
@@ -0,0 +1,397 @@
/* ******************************************************************************
*
*
* This program and the accompanying materials are made available under the
* terms of the Apache License, Version 2.0 which is available at
* https://www.apache.org/licenses/LICENSE-2.0.
*
* See the NOTICE file distributed with this work for additional
* information regarding copyright ownership.
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
* WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
* License for the specific language governing permissions and limitations
* under the License.
*
* SPDX-License-Identifier: Apache-2.0
******************************************************************************/
//
// @author Yurii Shyrma (iuriish@yahoo.com)
//
#include <ops/declarable/helpers/convolutions.h>
#include "cudnnUtils.h"
namespace sd {
namespace ops {
namespace platforms {
//////////////////////////////////////////////////////////////////////////
std::tuple<std::unique_ptr<NDArray>, std::unique_ptr<NDArray>> checkConv2dCUDNNPadAsymmetric(
NDArray* input, NDArray* gradI, const int iH, const int iW, const int oH, const int oW, const int kH,
const int kW, const int sH, const int sW, const int pH, const int pW, const int dH, const int dW,
const bool isNCHW) {
const auto pHsum = ((oH - 1) * sH + ((kH - 1) * dH + 1) - iH);
const auto pWsum = ((oW - 1) * sW + ((kW - 1) * dW + 1) - iW);
const bool isPHasymm = pH != (pHsum - pH);
const bool isPWasymm = pW != (pWsum - pW);
std::unique_ptr<NDArray> uNewInput = {}, uNewGradI = {};
if (!isPHasymm && !isPWasymm) return std::make_tuple(std::move(uNewInput), std::move(uNewGradI));
std::vector<LongType> newShape = input->getShapeAsVector();
const int iHposition = isNCHW ? 2 : 1;
if (isPHasymm) newShape[iHposition] += 1;
if (isPWasymm) newShape[iHposition + 1] += 1;
uNewInput.reset(new NDArray(input->ordering(), newShape, input->dataType(), input->getContext()));
if (isNCHW)
(*uNewInput)({0, 0, 0, 0, 0, input->sizeAt(2), 0, input->sizeAt(3)}).assign(input);
else
(*uNewInput)({0, 0, 0, input->sizeAt(1), 0, input->sizeAt(2), 0, 0}).assign(input);
if (gradI != nullptr)
uNewGradI.reset(new NDArray(gradI->ordering(), newShape, gradI->dataType(), gradI->getContext()));
return std::make_tuple(std::move(uNewInput), std::move(uNewGradI));
}
//////////////////////////////////////////////////////////////////////////
std::tuple<std::unique_ptr<NDArray>, std::unique_ptr<NDArray>> checkConv3dCUDNNPadAsymmetric(
NDArray* input, NDArray* gradI, const int iD, const int iH, const int iW, const int oD, const int oH,
const int oW, const int kD, const int kH, const int kW, const int sD, const int sH, const int sW, const int pD,
const int pH, const int pW, const int dD, const int dH, const int dW, const bool isNCDHW) {
const auto pDsum = ((oD - 1) * sD + ((kD - 1) * dD + 1) - iD);
const auto pHsum = ((oH - 1) * sH + ((kH - 1) * dH + 1) - iH);
const auto pWsum = ((oW - 1) * sW + ((kW - 1) * dW + 1) - iW);
const bool isPDasymm = pD != (pDsum - pD);
const bool isPHasymm = pH != (pHsum - pH);
const bool isPWasymm = pW != (pWsum - pW);
std::unique_ptr<NDArray> uNewInput = {}, uNewGradI = {};
if (!isPDasymm && !isPHasymm && !isPWasymm) return std::make_tuple(std::move(uNewInput), std::move(uNewGradI));
std::vector<LongType> newShape = input->getShapeAsVector();
const int iDposition = isNCDHW ? 2 : 1;
if (isPDasymm) newShape[iDposition] += 1;
if (isPHasymm) newShape[iDposition + 1] += 1;
if (isPWasymm) newShape[iDposition + 2] += 1;
uNewInput.reset(new NDArray(input->ordering(), newShape, input->dataType(), input->getContext()));
if (isNCDHW)
(*uNewInput)({0, 0, 0, 0, 0, input->sizeAt(2), 0, input->sizeAt(3), 0, input->sizeAt(4)}).assign(input);
else
(*uNewInput)({0, 0, 0, input->sizeAt(1), 0, input->sizeAt(2), 0, input->sizeAt(3), 0, 0}).assign(input);
if (gradI != nullptr)
uNewGradI.reset(new NDArray(gradI->ordering(), newShape, gradI->dataType(), gradI->getContext()));
return std::make_tuple(std::move(uNewInput), std::move(uNewGradI));
}
//////////////////////////////////////////////////////////////////////////
void pooling2dCUDNN(const LaunchContext* context, NDArray* input, NDArray* output, const int kH, const int kW,
const int sH, const int sW, const int pH, const int pW, const int dH, const int dW,
const bool isNCHW, const cudnnPoolingMode_t mode) {
LongType bS, iC, iH, iW, oC, oH,
oW; // batch size, input channels, input height/width, output channels, output height/width;
LongType indIOioC, indIiH, indWoC, indWiC, indWkH, indOoH; // corresponding indexes
ConvolutionUtils::getSizesAndIndexesConv2d(isNCHW, 0, *input, *output, bS, iC, iH, iW, oC, oH, oW, indIOioC, indIiH,
indWiC, indWoC, indWkH, indOoH);
auto handle = reinterpret_cast<cudnnHandle_t*>(context->getCuDnnHandle());
CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnSetStream), cudnnSetStream(*handle, *context->getCudaStream()));
cudnnTensorFormat_t format = isNCHW ? CUDNN_TENSOR_NCHW : CUDNN_TENSOR_NHWC;
// input descriptor, output descriptor
CudnnTensor x, z;
if (input->ordering() == 'c')
x.set4D(format, cudnnDataType(input->dataType()), bS, iC, iH, iW);
else
x.set4DEx(cudnnDataType(input->dataType()), bS, iC, iH, iW, input->strideAt(0), input->strideAt(indIOioC),
input->strideAt(indIiH), input->strideAt(indIiH + 1));
if (output->ordering() == 'c')
z.set4D(format, cudnnDataType(output->dataType()), bS, oC, oH, oW);
else
z.set4DEx(cudnnDataType(output->dataType()), bS, oC, oH, oW, output->strideAt(0), output->strideAt(indIOioC),
output->strideAt(indOoH), output->strideAt(indOoH + 1));
// description of pooling
PoolingDesc pooling;
pooling.set2D(mode, CUDNN_PROPAGATE_NAN, kH, kW, pH, pW, sH, sW);
// provide scaling parameters
const float alpha32(1), beta32(0);
const double alpha64(1), beta64(0);
const void* alpha =
output->sizeOfT() <= 4 ? reinterpret_cast<const void*>(&alpha32) : reinterpret_cast<const void*>(&alpha64);
const void* beta =
output->sizeOfT() <= 4 ? reinterpret_cast<const void*>(&beta32) : reinterpret_cast<const void*>(&beta64);
NDArray::prepareSpecialUse({output}, {input});
// run calculation
CHECK_CUDNN_FAILURE_MSG(
STRINGIZE(cudnnPoolingForward),
cudnnPoolingForward(*handle, pooling, alpha, x, input->specialBuffer(), beta, z, output->specialBuffer()));
auto cudaErr = cudaStreamSynchronize(*context->getCudaStream());
if (cudaErr != 0) throw cuda_exception::build("pooling2dCUDNN: cudaStreamSynchronize failed !", cudaErr);
NDArray::registerSpecialUse({output}, {input});
}
//////////////////////////////////////////////////////////////////////////
void pooling2dBpCUDNN(const LaunchContext* context, NDArray* input, NDArray* gradO, NDArray* gradI,
const int kH, const int kW, const int sH, const int sW, const int pH, const int pW, const int dH,
const int dW, const bool isNCHW, const cudnnPoolingMode_t mode) {
LongType bS, iC, iH, iW, oC, oH,
oW; // batch size, input channels, input height/width, output channels, output height/width;
LongType indIOioC, indIiH, indWoC, indWiC, indWkH, indOoH; // corresponding indexes
ConvolutionUtils::getSizesAndIndexesConv2d(isNCHW, 0, *input, *gradO, bS, iC, iH, iW, oC, oH, oW, indIOioC, indIiH,
indWiC, indWoC, indWkH, indOoH);
auto handle = reinterpret_cast<cudnnHandle_t*>(context->getCuDnnHandle());
CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnSetStream), cudnnSetStream(*handle, *context->getCudaStream()));
cudnnTensorFormat_t format = isNCHW ? CUDNN_TENSOR_NCHW : CUDNN_TENSOR_NHWC;
// input and gradI descriptor
CudnnTensor x;
if (input->ordering() == 'c')
x.set4D(format, cudnnDataType(input->dataType()), bS, iC, iH, iW);
else
x.set4DEx(cudnnDataType(input->dataType()), bS, iC, iH, iW, input->strideAt(0), input->strideAt(indIOioC),
input->strideAt(indIiH), input->strideAt(indIiH + 1));
// gradO descriptor
CudnnTensor dz;
if (gradO->ordering() == 'c')
dz.set4D(format, cudnnDataType(gradO->dataType()), bS, oC, oH, oW);
else
dz.set4DEx(cudnnDataType(gradO->dataType()), bS, oC, oH, oW, gradO->strideAt(0), gradO->strideAt(indIOioC),
gradO->strideAt(indOoH), gradO->strideAt(indOoH + 1));
// description of pooling
PoolingDesc pooling;
pooling.set2D(mode, CUDNN_PROPAGATE_NAN, kH, kW, pH, pW, sH, sW);
// provide scaling parameters
const float alpha32(1), beta32(0);
const double alpha64(1), beta64(0);
const void* alpha =
gradO->sizeOfT() <= 4 ? reinterpret_cast<const void*>(&alpha32) : reinterpret_cast<const void*>(&alpha64);
const void* beta =
gradO->sizeOfT() <= 4 ? reinterpret_cast<const void*>(&beta32) : reinterpret_cast<const void*>(&beta64);
NDArray::prepareSpecialUse({gradI}, {input, gradO});
// run calculation for gradI
CHECK_CUDNN_FAILURE_MSG(
STRINGIZE(cudnnPoolingBackward),
cudnnPoolingBackward(*handle, pooling, alpha, dz, gradO->specialBuffer(), dz, gradO->specialBuffer(), x,
input->specialBuffer(), beta, x, gradI->specialBuffer()));
auto cudaErr = cudaStreamSynchronize(*context->getCudaStream());
if (cudaErr != 0) throw cuda_exception::build("pooling2dBpCUDNN: cudaStreamSynchronize failed !", cudaErr);
NDArray::registerSpecialUse({gradI}, {input, gradO});
}
//////////////////////////////////////////////////////////////////////////
void pooling3dCUDNN(const LaunchContext* context, NDArray* input, NDArray* output, const int kD, const int kH,
const int kW, const int sD, const int sH, const int sW, const int pD, const int pH, const int pW,
const int dD, const int dH, const int dW, const bool isNCDHW, const cudnnPoolingMode_t mode) {
auto handle = reinterpret_cast<cudnnHandle_t*>(context->getCuDnnHandle());
CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnSetStream), cudnnSetStream(*handle, *context->getCudaStream()));
const int numDims = 5;
LongType bS, iC, iD, iH, iW, oC, oD, oH,
oW; // batch size, input channels, input depth/height/width, output channels, output depth/height/width;
LongType indIOioC, indIOioD, indWoC, indWiC, indWkD; // corresponding indexes
ConvolutionUtils::getSizesAndIndexesConv3d(isNCDHW, 0, *input, *output, bS, iC, iD, iH, iW, oC, oD, oH, oW, indIOioC,
indIOioD, indWiC, indWoC, indWkD);
const int pSizes[] = {pD, pH, pW};
const int sSizes[] = {sD, sH, sW};
const int kSizes[] = {kD, kH, kW};
const LongType xShape[] = {bS, iC, iD, iH, iW};
const LongType zShape[] = {bS, oC, oD, oH, oW};
const LongType xStrides[] = {(LongType)input->strideAt(0), (LongType)input->strideAt(1), (LongType)input->strideAt(2),
(LongType)input->strideAt(3), (LongType)input->strideAt(4)};
const LongType zStrides[] = {(LongType)output->strideAt(0), (LongType)output->strideAt(1), (LongType)output->strideAt(2),
(LongType)output->strideAt(3), (LongType)output->strideAt(4)};
cudnnTensorFormat_t format = isNCDHW ? CUDNN_TENSOR_NCHW : CUDNN_TENSOR_NHWC;
// input descriptor, output descriptor
CudnnTensor x, z;
if (input->ordering() == 'c') {
int newShape[5];
for(int i = 0; i < 5; i++) {
newShape[i] = static_cast<int>(xShape[i]);
}
x.setEx(format, cudnnDataType(input->dataType()), numDims, newShape);
} else {
int newShape[5];
int newStrides[5];
for(int i = 0; i < 5; i++) {
newShape[i] = static_cast<int>(xShape[i]);
}
for(int i = 0; i < 5; i++) {
newStrides[i] = static_cast<int>(xStrides[i]);
}
x.set(cudnnDataType(input->dataType()), numDims, newShape, newStrides);
}
if (output->ordering() == 'c') {
int newShape[5];
int newStrides[5];
for(int i = 0; i < 5; i++) {
newShape[i] = static_cast<int>(zShape[i]);
}
for(int i = 0; i < 5; i++) {
newStrides[i] = static_cast<int>(zStrides[i]);
}
z.setEx(format, cudnnDataType(output->dataType()), numDims, newShape);
} else {
int newShape[5];
int newStrides[5];
for(int i = 0; i < 5; i++) {
newShape[i] = static_cast<int>(zShape[i]);
}
for(int i = 0; i < 5; i++) {
newStrides[i] = static_cast<int>(zStrides[i]);
}
z.set(cudnnDataType(output->dataType()), numDims, newShape, newStrides);
}
// description of pooling
PoolingDesc pooling;
pooling.set(mode, CUDNN_PROPAGATE_NAN, numDims - 2, kSizes, pSizes, sSizes);
// provide scaling parameters
const float alpha32(1), beta32(0);
const double alpha64(1), beta64(0);
const void* alpha =
output->sizeOfT() <= 4 ? reinterpret_cast<const void*>(&alpha32) : reinterpret_cast<const void*>(&alpha64);
const void* beta =
output->sizeOfT() <= 4 ? reinterpret_cast<const void*>(&beta32) : reinterpret_cast<const void*>(&beta64);
NDArray::prepareSpecialUse({output}, {input});
// run calculation
CHECK_CUDNN_FAILURE_MSG(
STRINGIZE(cudnnPoolingForward),
cudnnPoolingForward(*handle, pooling, alpha, x, input->specialBuffer(), beta, z, output->specialBuffer()));
auto cudaErr = cudaStreamSynchronize(*context->getCudaStream());
if (cudaErr != 0) throw cuda_exception::build("pooling3dCUDNN: cudaStreamSynchronize failed !", cudaErr);
NDArray::registerSpecialUse({output}, {input});
}
//////////////////////////////////////////////////////////////////////////
void pooling3dBpCUDNN(const LaunchContext* context, NDArray* input, NDArray* gradO, NDArray* gradI,
const int kD, const int kH, const int kW, const int sD, const int sH, const int sW, const int pD,
const int pH, const int pW, const int dD, const int dH, const int dW, const bool isNCDHW,
const cudnnPoolingMode_t mode) {
auto handle = reinterpret_cast<cudnnHandle_t*>(context->getCuDnnHandle());
CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnSetStream), cudnnSetStream(*handle, *context->getCudaStream()));
const int numDims = 5;
LongType bS, iC, iD, iH, iW, oC, oD, oH,
oW; // batch size, input channels, input depth/height/width, output channels, output depth/height/width;
LongType indIOioC, indIOioD, indWoC, indWiC, indWkD; // corresponding indexes
ConvolutionUtils::getSizesAndIndexesConv3d(isNCDHW, 0, *input, *gradO, bS, iC, iD, iH, iW, oC, oD, oH, oW, indIOioC,
indIOioD, indWiC, indWoC, indWkD);
const int pSizes[] = {pD, pH, pW};
const int sSizes[] = {sD, sH, sW};
const int kSizes[] = {kD, kH, kW};
const int xShape[] = {(int) bS, (int) iC, (int) iD, (int) iH, (int) iW};
const int dzShape[] = {(int) bS, (int) oC, (int) oD, (int) oH,(int) oW};
const int xStrides[] = { (int) input->strideAt(0), (int)input->strideAt(1), (int)input->strideAt(2),
(int)input->strideAt(3), (int)input->strideAt(4)};
const int dzStrides[] = {(int)gradO->strideAt(0), (int)gradO->strideAt(1), (int)gradO->strideAt(2),
(int)gradO->strideAt(3), (int)gradO->strideAt(4)};
cudnnTensorFormat_t format = isNCDHW ? CUDNN_TENSOR_NCHW : CUDNN_TENSOR_NHWC;
// input and gradI descriptor
CudnnTensor x;
if ( input->ordering() == 'c') {
x.setEx(format, cudnnDataType(input->dataType()), numDims, xShape);
} else {
x.set(cudnnDataType(input->dataType()), numDims, xShape, xStrides);
}
// gradO descriptor
CudnnTensor dz;
if ( gradO->ordering() == 'c')
dz.setEx(format, cudnnDataType(gradO->dataType()), numDims, dzShape);
else
dz.set(cudnnDataType(gradO->dataType()), numDims, dzShape, dzStrides);
// description of pooling
PoolingDesc pooling;
pooling.set(mode, CUDNN_PROPAGATE_NAN, numDims - 2, kSizes, pSizes, sSizes);
// provide scaling parameters
const float alpha32(1), beta32(0);
const double alpha64(1), beta64(0);
const void* alpha =
gradO->sizeOfT() <= 4 ? reinterpret_cast<const void*>(&alpha32) : reinterpret_cast<const void*>(&alpha64);
const void* beta =
gradO->sizeOfT() <= 4 ? reinterpret_cast<const void*>(&beta32) : reinterpret_cast<const void*>(&beta64);
// cudnn maxpool2d_bp api requires ff output as one of input arguments
if (mode == CUDNN_POOLING_MAX) {
NDArray temp(gradO);
NDArray::prepareSpecialUse({gradI}, {input, gradO, &temp});
// run ff calculation
CHECK_CUDNN_FAILURE_MSG(
STRINGIZE(cudnnPoolingForward),
cudnnPoolingForward(*handle, pooling, alpha, x, input->specialBuffer(), beta, dz, temp.specialBuffer()));
// run bp calculation for gradI
CHECK_CUDNN_FAILURE_MSG(
STRINGIZE(cudnnPoolingBackward),
cudnnPoolingBackward(*handle, pooling, alpha, dz, temp.specialBuffer(), dz, gradO->specialBuffer(), x,
input->specialBuffer(), beta, x, gradI->specialBuffer()));
NDArray::registerSpecialUse({gradI}, {input, gradO, &temp});
} else {
NDArray::prepareSpecialUse({gradI}, {input, gradO});
// run bp calculation for gradI
CHECK_CUDNN_FAILURE_MSG(
STRINGIZE(cudnnPoolingBackward),
cudnnPoolingBackward(*handle, pooling, alpha, dz, gradO->specialBuffer(), dz, gradO->specialBuffer(), x,
input->specialBuffer(), beta, x, gradI->specialBuffer()));
NDArray::registerSpecialUse({gradI}, {input, gradO});
}
auto cudaErr = cudaStreamSynchronize(*context->getCudaStream());
if (cudaErr != 0) throw cuda_exception::build("pooling3dBpCUDNN: cudaStreamSynchronize failed !", cudaErr);
}
} // namespace platforms
} // namespace ops
} // namespace sd
libnd4j/include/ops/declarable/platform/cudnn/cudnnUtils.h
+358
View File
@@ -0,0 +1,358 @@
/* ******************************************************************************
*
*
* This program and the accompanying materials are made available under the
* terms of the Apache License, Version 2.0 which is available at
* https://www.apache.org/licenses/LICENSE-2.0.
*
* See the NOTICE file distributed with this work for additional
* information regarding copyright ownership.
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
* WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
* License for the specific language governing permissions and limitations
* under the License.
*
* SPDX-License-Identifier: Apache-2.0
******************************************************************************/
//
// @author raver119@gmail.com
//
#ifndef SD_CUDNNUTILS_H
#define SD_CUDNNUTILS_H
#include <cudnn.h>
#include <exceptions/cuda_exception.h>
#include <exceptions/datatype_exception.h>
#include <helpers/PointersManager.h>
#include <ops/declarable/OpRegistrator.h>
#include <ops/declarable/PlatformHelper.h>
#include <system/platform_boilerplate.h>
#include <memory>
#include <tuple>
#include <vector>
#define CUDNN_NEW_RNN_API_VER 8001
#define CUDNN_CLIPPING_API_VER 7201
namespace sd {
namespace ops {
namespace platforms {
DECLARE_PLATFORM(conv2d, ENGINE_CUDA);
DECLARE_PLATFORM(conv2d_bp, ENGINE_CUDA);
DECLARE_PLATFORM(conv3dnew, ENGINE_CUDA);
DECLARE_PLATFORM(conv3dnew_bp, ENGINE_CUDA);
DECLARE_PLATFORM(depthwise_conv2d, ENGINE_CUDA);
DECLARE_PLATFORM(depthwise_conv2d_bp, ENGINE_CUDA);
DECLARE_PLATFORM(batchnorm, ENGINE_CUDA);
DECLARE_PLATFORM(batchnorm_bp, ENGINE_CUDA);
DECLARE_PLATFORM(avgpool2d, ENGINE_CUDA);
DECLARE_PLATFORM(avgpool2d_bp, ENGINE_CUDA);
DECLARE_PLATFORM(maxpool2d, ENGINE_CUDA);
DECLARE_PLATFORM(maxpool2d_bp, ENGINE_CUDA);
DECLARE_PLATFORM(avgpool3dnew, ENGINE_CUDA);
DECLARE_PLATFORM(avgpool3dnew_bp, ENGINE_CUDA);
DECLARE_PLATFORM(maxpool3dnew, ENGINE_CUDA);
DECLARE_PLATFORM(maxpool3dnew_bp, ENGINE_CUDA);
DECLARE_PLATFORM(lstmLayer, ENGINE_CUDA);
DECLARE_PLATFORM(ctc_loss, ENGINE_CUDA);
DECLARE_PLATFORM(ctc_loss_grad, ENGINE_CUDA);
//////////////////////////////////////////////////////////////////////////
inline void throwIfCudnnFailed(cudnnStatus_t result_status,
const char* message = "Cudnn error: ", const char* prefix = nullptr) {
if (result_status != CUDNN_STATUS_SUCCESS) {
std::string err_message;
if (prefix) err_message = std::string(prefix) + ": ";
err_message += std::string(message);
throw cuda_exception::build(err_message, result_status);
}
}
#define STRINGIZE(x) STRINGIZE2(x)
#define STRINGIZE2(x) #x
#define CHECK_CUDNN_FAILURE(result_status) throwIfCudnnFailed(result_status, "")
#define CHECK_CUDNN_FAILURE_MSG(custom_message, result_status) \
throwIfCudnnFailed(result_status, custom_message, __func__)
template <typename T>
SD_INLINE const T* bufferInHost(NDArray& array) {
array.syncToHost();
return reinterpret_cast<const T*>(array.buffer());
}
#define MOVEONLY_DESC_IMPL(DESC) \
DESC(const DESC& s) = delete; \
DESC& operator=(const DESC& other) = delete; \
DESC(DESC&& other) noexcept : desc(std::move(other.desc)) { other.desc = {}; } \
DESC& operator=(DESC&& other) noexcept { \
if (&other == this) return *this; \
destroy(); \
desc = std::move(other.desc); \
other.desc = {}; \
return *this; \
}
#define MOVEONLY_DESC_FULL_IMPL(DESC_CLASS, DESC_NAME) \
DESC_CLASS() { create(); } \
DESC_CLASS(cudnn##DESC_NAME##_t created) { desc = created; } \
void create() { CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnCreate##DESC_NAME), cudnnCreate##DESC_NAME(&desc)); } \
void destroy() { \
if (desc) CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnCreate##DESC_NAME), cudnnDestroy##DESC_NAME(desc)); \
desc = {}; \
} \
MOVEONLY_DESC_IMPL(DESC_CLASS) \
operator cudnn##DESC_NAME##_t() const { return desc; } \
~DESC_CLASS() { destroy(); } \
cudnn##DESC_NAME##_t desc;
struct CudnnTensor {
MOVEONLY_DESC_FULL_IMPL(CudnnTensor, TensorDescriptor)
template <typename... Args>
void set(Args&&... args) {
CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnSetTensorNdDescriptor),
cudnnSetTensorNdDescriptor(desc, std::forward<Args>(args)...));
}
template <typename... Args>
void setEx(Args&&... args) {
CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnSetTensorNdDescriptorEx),
cudnnSetTensorNdDescriptorEx(desc, std::forward<Args>(args)...));
}
template <typename... Args>
void set4D(Args&&... args) {
CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnSetTensor4dDescriptor),
cudnnSetTensor4dDescriptor(desc, std::forward<Args>(args)...));
}
template <typename... Args>
void set4DEx(Args&&... args) {
CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnSetTensor4dDescriptorEx),
cudnnSetTensor4dDescriptorEx(desc, std::forward<Args>(args)...));
}
};
struct CudnnTensorList {
MOVEONLY_DESC_IMPL(CudnnTensorList)
CudnnTensorList(int size) {
desc.resize(size);
for (int i = 0; i < size; i++) {
CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnCreateTensorDescriptor), cudnnCreateTensorDescriptor(&desc[i]));
}
}
template <typename... Args>
void set(int index, Args&&... args) {
if (index < desc.size()) {
CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnSetTensorNdDescriptor),
cudnnSetTensorNdDescriptor(desc[index], std::forward<Args>(args)...));
}
}
cudnnTensorDescriptor_t get(int i) const {
if (i < desc.size()) return desc[i];
return nullptr;
}
const cudnnTensorDescriptor_t* getDescriptors() const { return desc.data(); }
void destroy() {
for (auto x : desc) {
CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnDestroyTensorDescriptor), cudnnDestroyTensorDescriptor(x));
}
desc = {};
}
~CudnnTensorList() { destroy(); }
std::vector<cudnnTensorDescriptor_t> desc;
};
struct FilterDesc {
MOVEONLY_DESC_FULL_IMPL(FilterDesc, FilterDescriptor)
template <typename... Args>
void set(Args&&... args) {
CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnSetFilterNdDescriptor),
cudnnSetFilterNdDescriptor(desc, std::forward<Args>(args)...));
}
template <typename... Args>
void set4D(Args&&... args) {
CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnSetFilter4dDescriptor),
cudnnSetFilter4dDescriptor(desc, std::forward<Args>(args)...));
}
};
struct DropoutDesc {
MOVEONLY_DESC_FULL_IMPL(DropoutDesc, DropoutDescriptor)
template <typename... Args>
void set(Args&&... args) {
CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnSetDropoutDescriptor),
cudnnSetDropoutDescriptor(desc, std::forward<Args>(args)...));
}
};
#if CUDNN_VERSION > CUDNN_NEW_RNN_API_VER
struct RnnDataDesc {
MOVEONLY_DESC_FULL_IMPL(RnnDataDesc, RNNDataDescriptor)
template <typename... Args>
void set(Args&&... args) {
CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnSetRNNDataDescriptor),
cudnnSetRNNDataDescriptor(desc, std::forward<Args>(args)...));
}
};
#endif
SD_INLINE void setRnnDescriptorOldApi(cudnnRNNDescriptor_t rnnDesc, cudnnHandle_t handle, cudnnRNNInputMode_t inputMode,
cudnnDirectionMode_t dirMode, cudnnRNNMode_t cellMode, cudnnRNNAlgo_t algo,
cudnnDataType_t mathPrec, int32_t hiddenSize, int32_t numLayers,
cudnnDropoutDescriptor_t dropoutDesc, bool use_tensor_op = false) {
auto err = cudnnSetRNNDescriptor_v6(handle, rnnDesc, hiddenSize, numLayers, dropoutDesc, inputMode, dirMode, cellMode,
algo, mathPrec);
CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnSetRNNDescriptor_v6), err);
#if CUDNN_VERSION >= 7001
if (cudnnGetVersion() >= 7001) {
cudnnMathType_t mathType = use_tensor_op ? CUDNN_TENSOR_OP_MATH : CUDNN_DEFAULT_MATH;
CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnSetRNNMatrixMathType), cudnnSetRNNMatrixMathType(rnnDesc, mathType));
}
#endif
return;
}
struct RnnDesc {
MOVEONLY_DESC_FULL_IMPL(RnnDesc, RNNDescriptor)
template <typename... Args>
void setUsingOldAPI(Args&&... args) {
setRnnDescriptorOldApi(desc, std::forward<Args>(args)...);
}
#if CUDNN_VERSION >= CUDNN_NEW_RNN_API_VER
template <typename... Args>
void set(Args&&... args) {
CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnSetRNNDescriptor_v8),
cudnnSetRNNDescriptor_v8(desc, std::forward<Args>(args)...));
}
#endif
};
struct CTCLossDesc {
MOVEONLY_DESC_FULL_IMPL(CTCLossDesc, CTCLossDescriptor)
template <typename... Args>
void set(Args&&... args) {
CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnSetCTCLossDescriptorEx),
cudnnSetCTCLossDescriptorEx(desc, std::forward<Args>(args)...));
}
};
//////////////////////////////////////////////////////////////////////////
struct PoolingDesc {
MOVEONLY_DESC_FULL_IMPL(PoolingDesc, PoolingDescriptor)
template <typename... Args>
void set(Args&&... args) {
CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnSetPoolingNdDescriptor),
cudnnSetPoolingNdDescriptor(desc, std::forward<Args>(args)...));
}
template <typename... Args>
void set2D(Args&&... args) {
CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnSetPooling2dDescriptor),
cudnnSetPooling2dDescriptor(desc, std::forward<Args>(args)...));
}
};
//////////////////////////////////////////////////////////////////////////
struct ConvolutionDesc {
MOVEONLY_DESC_FULL_IMPL(ConvolutionDesc, ConvolutionDescriptor)
template <typename... Args>
void set(Args&&... args) {
CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnSetConvolutionNdDescriptor),
cudnnSetConvolutionNdDescriptor(desc, std::forward<Args>(args)...));
}
template <typename... Args>
void set2D(Args&&... args) {
CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnSetConvolution2dDescriptor),
cudnnSetConvolution2dDescriptor(desc, std::forward<Args>(args)...));
}
};
//////////////////////////////////////////////////////////////////////////
SD_INLINE cudnnDataType_t cudnnDataType(DataType dataType) {
switch (dataType) {
case FLOAT32:
return CUDNN_DATA_FLOAT;
case DOUBLE:
return CUDNN_DATA_DOUBLE;
case HALF:
return CUDNN_DATA_HALF;
case INT32:
return CUDNN_DATA_INT32;
case INT8:
return CUDNN_DATA_INT8;
default:
throw datatype_exception::build("Unsupported data type", dataType);
}
}
//////////////////////////////////////////////////////////////////////////
std::tuple<std::unique_ptr<NDArray>, std::unique_ptr<NDArray>> checkConv2dCUDNNPadAsymmetric(
NDArray* input, NDArray* gradI, const int iH, const int iW, const int oH, const int oW, const int kH,
const int kW, const int sH, const int sW, const int pH, const int pW, const int dH, const int dW,
const bool isNCHW);
//////////////////////////////////////////////////////////////////////////
std::tuple<std::unique_ptr<NDArray>, std::unique_ptr<NDArray>> checkConv3dCUDNNPadAsymmetric(
NDArray* input, NDArray* gradI, const int iD, const int iH, const int iW, const int oD, const int oH,
const int oW, const int kD, const int kH, const int kW, const int sD, const int sH, const int sW, const int pD,
const int pH, const int pW, const int dD, const int dH, const int dW, const bool isNCDHW);
//////////////////////////////////////////////////////////////////////////
void pooling2dCUDNN(const LaunchContext* context, NDArray* input, NDArray* output, const int kH, const int kW,
const int sH, const int sW, const int pH, const int pW, const int dH, const int dW,
const bool isNCHW, const cudnnPoolingMode_t mode);
//////////////////////////////////////////////////////////////////////////
void pooling2dBpCUDNN(const LaunchContext* context, NDArray* input, NDArray* gradO, NDArray* gradI,
const int kH, const int kW, const int sH, const int sW, const int pH, const int pW, const int dH,
const int dW, const bool isNCHW, const cudnnPoolingMode_t mode);
//////////////////////////////////////////////////////////////////////////
void pooling3dCUDNN(const LaunchContext* context, NDArray* input, NDArray* output, const int kD, const int kH,
const int kW, const int sD, const int sH, const int sW, const int pD, const int pH, const int pW,
const int dD, const int dH, const int dW, const bool isNCDHW, const cudnnPoolingMode_t mode);
//////////////////////////////////////////////////////////////////////////
void pooling3dBpCUDNN(const LaunchContext* context, NDArray* input, NDArray* gradO, NDArray* gradI,
const int kD, const int kH, const int kW, const int sD, const int sH, const int sW, const int pD,
const int pH, const int pW, const int dD, const int dH, const int dW, const bool isNCDHW,
const cudnnPoolingMode_t mode);
} // namespace platforms
} // namespace ops
} // namespace sd
#endif // SD_CUDNNUTILS_H
libnd4j/include/ops/declarable/platform/cudnn/depthwiseConv2d.cu
+533
View File
@@ -0,0 +1,533 @@
/* ******************************************************************************
*
*
* This program and the accompanying materials are made available under the
* terms of the Apache License, Version 2.0 which is available at
* https://www.apache.org/licenses/LICENSE-2.0.
*
* See the NOTICE file distributed with this work for additional
* information regarding copyright ownership.
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
* WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
* License for the specific language governing permissions and limitations
* under the License.
*
* SPDX-License-Identifier: Apache-2.0
******************************************************************************/
//
// @author Yurii Shyrma (iuriish@yahoo.com)
//
#include <ops/declarable/helpers/convolutions.h>
#include "cudnnUtils.h"
namespace sd {
namespace ops {
namespace platforms {
//////////////////////////////////////////////////////////////////////////
static void depthwiseConv2dCUDNN(const LaunchContext* context, NDArray* input, NDArray* weights,
NDArray* bias, NDArray* output, const LongType kH, const LongType kW, const LongType sH,
const LongType sW, const LongType pH, const LongType pW, const LongType dH, const LongType dW,
const LongType paddingMode, const bool isNCHW) {
// cudnn supports only following case: mC = 1, oC = iC (groupCount == iC)
// input [bS, iC, iH, iW] nchw or [bS, iH, iW, iC] nhwc
// weights [iC, mC, kH, kW]
// bias [oC], may be nullptr
// output [bS, oC, oH, oW] nchw or [bS, oH, oW, oC] nhwc
// oC = iC*mC
LongType bS, iC, iH, iW, mC, oC, oH,
oW; // batch size, input channels, input height/width, output channels, output height/width;
LongType indIOioC, indIiH, indWmC, indWiC, indWkH, indOoH; // corresponding indexes
ConvolutionUtils::getSizesAndIndexesConv2d(isNCHW, 0, *input, *output, bS, iC, iH, iW, oC, oH, oW, indIOioC, indIiH,
indWiC, indWmC, indWkH, indOoH);
mC = weights->sizeAt(1);
auto handle = reinterpret_cast<cudnnHandle_t*>(context->getCuDnnHandle());
CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnSetStream), cudnnSetStream(*handle, *context->getCudaStream()));
cudnnTensorFormat_t format = isNCHW ? CUDNN_TENSOR_NCHW : CUDNN_TENSOR_NHWC;
PointersManager manager(context, __func__);
// input descriptor
CudnnTensor x;
if (input->ordering() == 'c')
x.set4D(format, cudnnDataType(input->dataType()), bS, iC, iH, iW);
else
x.set4DEx(cudnnDataType(input->dataType()), bS, iC, iH, iW, input->strideAt(0), input->strideAt(indIOioC),
input->strideAt(indIiH), input->strideAt(indIiH + 1));
// weights descriptor
FilterDesc w;
w.set4D(cudnnDataType(weights->dataType()), CUDNN_TENSOR_NCHW, iC, mC, kH, kW);
// output descriptor
CudnnTensor z;
if (output->ordering() == 'c')
z.set4D(format, cudnnDataType(output->dataType()), bS, oC, oH, oW);
else
z.set4DEx(cudnnDataType(output->dataType()), bS, oC, oH, oW, output->strideAt(0), output->strideAt(indIOioC),
output->strideAt(indOoH), output->strideAt(indOoH + 1));
// description of convolution
ConvolutionDesc conv;
conv.set2D(pH, pW, sH, sW, dH, dW, CUDNN_CROSS_CORRELATION, cudnnDataType(output->dataType()));
CHECK_CUDNN_FAILURE_MSG(
STRINGIZE(cudnnSetConvolutionGroupCount),
cudnnSetConvolutionGroupCount(
conv, iC)); // set number of groups (depthwise mode) in description of convolution, groupCount == iC
// algorithm description
cudnnConvolutionFwdAlgo_t algo;
cudnnConvolutionFwdAlgoPerf_t algoPerf;
int count = 0;
// CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnGetConvolutionForwardAlgorithm), cudnnGetConvolutionForwardAlgorithm(
// *handle, x, w, conv, z, CUDNN_CONVOLUTION_FWD_PREFER_FASTEST, 0, &algo));
CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnFindConvolutionForwardAlgorithm),
cudnnFindConvolutionForwardAlgorithm(*handle, x, w, conv, z, 1, &count, &algoPerf));
if (count == 0)
throw cuda_exception::build("depthwiseConv2dCUDNN: cudnnGetConvolutionForwardAlgorithm failed", 0);
algo = algoPerf.algo;
// allocate auxiliary device memory, abbreviation ws means workspace
size_t wsSize;
CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnGetConvolutionForwardWorkspaceSize),
cudnnGetConvolutionForwardWorkspaceSize(*handle, x, w, conv, z, algo, &wsSize));
void* wsData = manager.allocateDevMem(wsSize);
// provide scaling parameters
const float alpha32(1), beta32(0);
const double alpha64(1), beta64(0);
const void* alpha =
output->sizeOfT() <= 4 ? reinterpret_cast<const void*>(&alpha32) : reinterpret_cast<const void*>(&alpha64);
const void* beta =
output->sizeOfT() <= 4 ? reinterpret_cast<const void*>(&beta32) : reinterpret_cast<const void*>(&beta64);
NDArray::prepareSpecialUse({output}, {input, weights, bias});
// run calculation
CHECK_CUDNN_FAILURE_MSG(
STRINGIZE(cudnnConvolutionForward),
cudnnConvolutionForward(*handle, alpha, x, input->specialBuffer(), w, weights->specialBuffer(), conv, algo,
wsData, wsSize, beta, z, output->specialBuffer()));
// add bias if it is present
if (bias != nullptr) {
CudnnTensor b;
// b.set( format, cudnnDataType(bias->dataType()), 1, isNCHW ? bias->lengthOf() : 1, 1, isNCHW ? 1:
// bias->lengthOf());
b.set4D(CUDNN_TENSOR_NCHW, cudnnDataType(bias->dataType()), 1, oC, 1, 1);
CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnAddTensor), cudnnAddTensor(*handle, alpha, b, bias->specialBuffer(), alpha,
z, output->specialBuffer()));
}
}
//////////////////////////////////////////////////////////////////////////
static void depthwiseConv2dBpCUDNN(const LaunchContext* context, NDArray* input, NDArray* weights,
NDArray* gradO, NDArray* gradI, NDArray* gradW, NDArray* gradB, const LongType kH,
const LongType kW, const LongType sH, const LongType sW, const LongType pH, const LongType pW, const LongType dH,
const LongType dW, const LongType paddingMode, const bool isNCHW) {
// cudnn supports only following case: mC = 1, oC = iC (groupCount == iC)
// input, gradI [bS, iC, iH, iW] nchw or [bS, iH, iW, iC] nhwc
// weights, gradW [iC, mC, kH, kW]
// gradB [oC], may be nullptr
// gradO [bS, oC, oH, oW] nchw or [bS, oH, oW, oC] nhwc
// oC = iC*mC
LongType bS, iC, iH, iW, mC, oC, oH,
oW; // batch size, input channels, input height/width, output channels, output height/width;
LongType indIOioC, indIiH, indWmC, indWiC, indWkH, indOoH; // corresponding indexes
ConvolutionUtils::getSizesAndIndexesConv2d(isNCHW, 0, *input, *gradO, bS, iC, iH, iW, oC, oH, oW, indIOioC, indIiH,
indWiC, indWmC, indWkH, indOoH);
mC = weights->sizeAt(1);
auto handle = reinterpret_cast<cudnnHandle_t*>(context->getCuDnnHandle());
CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnSetStream), cudnnSetStream(*handle, *context->getCudaStream()));
cudnnTensorFormat_t format = isNCHW ? CUDNN_TENSOR_NCHW : CUDNN_TENSOR_NHWC;
PointersManager manager(context, __func__);
// input descriptor
CudnnTensor x;
if (input->ordering() == 'c')
x.set4D(format, cudnnDataType(input->dataType()), bS, iC, iH, iW);
else
x.set4DEx(cudnnDataType(input->dataType()), bS, iC, iH, iW, input->strideAt(0), input->strideAt(indIOioC),
input->strideAt(indIiH), input->strideAt(indIiH + 1));
// gradO descriptor
CudnnTensor dz;
if (gradO->ordering() == 'c')
dz.set4D(format, cudnnDataType(gradO->dataType()), bS, oC, oH, oW);
else
dz.set4DEx(cudnnDataType(gradO->dataType()), bS, oC, oH, oW, gradO->strideAt(0), gradO->strideAt(indIOioC),
gradO->strideAt(indOoH), gradO->strideAt(indOoH + 1));
// gradI descriptor
CudnnTensor dx;
if (gradI->ordering() == 'c')
dx.set4D(format, cudnnDataType(gradI->dataType()), bS, iC, iH, iW);
else
dx.set4DEx(cudnnDataType(gradI->dataType()), bS, iC, iH, iW, gradI->strideAt(0), gradI->strideAt(indIOioC),
gradI->strideAt(indIiH), gradI->strideAt(indIiH + 1));
// gradW descriptor
FilterDesc dw;
dw.set4D(cudnnDataType(gradW->dataType()), CUDNN_TENSOR_NCHW, iC, mC, kH, kW);
// description of convolution
ConvolutionDesc conv;
conv.set2D(pH, pW, sH, sW, dH, dW, CUDNN_CROSS_CORRELATION, cudnnDataType(gradO->dataType()));
CHECK_CUDNN_FAILURE_MSG(
STRINGIZE(cudnnSetConvolutionGroupCount),
cudnnSetConvolutionGroupCount(
conv, iC)); // set number of groups (depthwise mode) in description of convolution, groupCount == iC
// gradW algorithm description
cudnnConvolutionBwdFilterAlgo_t algoGradW;
cudnnConvolutionBwdFilterAlgoPerf_t algoGradWPerf;
int count = 0;
// CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnGetConvolutionBackwardFilterAlgorithm),
// cudnnGetConvolutionBackwardFilterAlgorithm( *handle, x, dz, conv, dw, CUDNN_CONVOLUTION_BWD_FILTER_PREFER_FASTEST,
// 0, &algoGradW));
CHECK_CUDNN_FAILURE_MSG(
STRINGIZE(cudnnFindConvolutionBackwardFilterAlgorithm),
cudnnFindConvolutionBackwardFilterAlgorithm(*handle, x, dz, conv, dw, 1, &count, &algoGradWPerf));
if (count == 0)
throw cuda_exception::build(
"depthwiseConv2dBpCUDNN: cudnnGetConvolutionBackwardFilterAlgorithm failed as the count is 0 ", 0);
algoGradW = algoGradWPerf.algo;
// gradI algorithm description
cudnnConvolutionBwdDataAlgo_t algoGradI;
cudnnConvolutionBwdDataAlgoPerf_t algoGradIPerf;
CHECK_CUDNN_FAILURE_MSG(
STRINGIZE(cudnnFindConvolutionBackwardDataAlgorithm),
cudnnFindConvolutionBackwardDataAlgorithm(*handle, dw, dz, conv, x, 1, &count, &algoGradIPerf));
if (count == 0)
throw cuda_exception::build(
"depthwiseConv2dBpCUDNN: cudnnGetConvolutionBackwardDataAlgorithm failed as the count is 0 ", 0);
algoGradI = algoGradIPerf.algo;
// allocate auxiliary device memory for gradW calculation, abbreviation ws means workspace
size_t wsGradWSize;
CHECK_CUDNN_FAILURE_MSG(
STRINGIZE(cudnnGetConvolutionBackwardFilterWorkspaceSize),
cudnnGetConvolutionBackwardFilterWorkspaceSize(*handle, x, dz, conv, dw, algoGradW, &wsGradWSize));
void* wsGradWData = manager.allocateDevMem(wsGradWSize);
// allocate auxiliary device memory for gradI calculation, abbreviation ws means workspace
size_t wsGradISize;
CHECK_CUDNN_FAILURE_MSG(
STRINGIZE(cudnnGetConvolutionBackwardDataWorkspaceSize),
cudnnGetConvolutionBackwardDataWorkspaceSize(*handle, dw, dz, conv, dx, algoGradI, &wsGradISize));
void* wsGradIData = manager.allocateDevMem(wsGradISize);
// provide scaling parameters
const float alpha32(1), beta32(0);
const double alpha64(1), beta64(0);
const void* alpha =
gradO->sizeOfT() <= 4 ? reinterpret_cast<const void*>(&alpha32) : reinterpret_cast<const void*>(&alpha64);
const void* beta =
gradO->sizeOfT() <= 4 ? reinterpret_cast<const void*>(&beta32) : reinterpret_cast<const void*>(&beta64);
NDArray::prepareSpecialUse({gradI, gradW, gradB}, {input, weights, gradO});
// run calculation for gradB (if not nullptr)
if (gradB != nullptr) {
CudnnTensor db;
// db.set( format, cudnnDataType(gradB->dataType()), 1, isNCHW ? gradB->lengthOf() : 1, 1, isNCHW ? 1:
// gradB->lengthOf());
db.set4D(CUDNN_TENSOR_NCHW, cudnnDataType(gradB->dataType()), 1, oC, 1, 1);
CHECK_CUDNN_FAILURE_MSG(
STRINGIZE(cudnnConvolutionBackwardBias),
cudnnConvolutionBackwardBias(*handle, alpha, dz, gradO->specialBuffer(), beta, db, gradB->specialBuffer()));
}
// run calculation for gradW
CHECK_CUDNN_FAILURE_MSG(
STRINGIZE(cudnnConvolutionBackwardFilter),
cudnnConvolutionBackwardFilter(*handle, alpha, x, input->specialBuffer(), dz, gradO->specialBuffer(), conv,
algoGradW, wsGradWData, wsGradWSize, beta, dw, gradW->specialBuffer()));
// run calculation for gradI
CHECK_CUDNN_FAILURE_MSG(
STRINGIZE(cudnnConvolutionBackwardData),
cudnnConvolutionBackwardData(*handle, alpha, dw, weights->specialBuffer(), dz, gradO->specialBuffer(), conv,
algoGradI, wsGradIData, wsGradISize, beta, dx, gradI->specialBuffer()));
NDArray::registerSpecialUse({gradI, gradW, gradB}, {input, weights, gradO});
}
//////////////////////////////////////////////////////////////////////////
PLATFORM_IMPL(depthwise_conv2d, ENGINE_CUDA) {
auto input = INPUT_VARIABLE(0); // [bS, iH, iW, iC] (NHWC) or [bS, iC, iH, iW] (NCHW)
auto weights = INPUT_VARIABLE(1); // [kH, kW, iC, mC], [mC, iC, kH, kW], [mC, kH, kW, iC]
auto bias = block.width() > 2 ? INPUT_VARIABLE(2) : nullptr; // [oC] = iC*mC
auto output = OUTPUT_VARIABLE(0); // [bS, oH, oW, iC*mC] (NHWC) or [bS, iC*mC, oH, oW] (NCHW)
REQUIRE_TRUE(input->rankOf() == 4, 0,
"DEPTHWISECONV2D CUDNN OP: rank of input array must be equal to 4, but got %i instead !",
input->rankOf());
REQUIRE_TRUE(weights->rankOf() == 4, 0,
"DEPTHWISECONV2D CUDNN OP: rank of weights array must be equal to 4, but got %i instead !",
weights->rankOf());
LongType kH = INT_ARG(0) > 0 ? INT_ARG(0) : static_cast<LongType>(weights->sizeAt(0)); // filter(kernel) height
LongType kW = INT_ARG(1) > 0 ? INT_ARG(1) : static_cast<LongType>(weights->sizeAt(1)); // filter(kernel) width
LongType sH = INT_ARG(2); // strides height
LongType sW = INT_ARG(3); // strides width
LongType pH = INT_ARG(4); // paddings height
LongType pW = INT_ARG(5); // paddings width
LongType dH = INT_ARG(6); // dilations height
LongType dW = INT_ARG(7); // dilations width
int paddingMode = INT_ARG(8); // 0-VALID, 1-SAME
int isNCHW = block.getIArguments()->size() > 9 ? !INT_ARG(9) : 1; // INT_ARG(9): 0-NCHW, 1-NHWC
int wFormat = block.getIArguments()->size() > 10
? INT_ARG(10)
: 0; // 0 - [kH, kW, iC, mC], 1 - [mC, iC, kH, kW], 2 - [mC, kH, kW, iC]
LongType bS, iC, iH, iW, mC, oC, oH, oW; // batch size, input channels, input height/width, channels multiplier(oC =
// iC*mC), output channels, output height/width
LongType indIOioC, indIiH, indWmC, indWiC, indWkH, indOoH; // corresponding indexes
ConvolutionUtils::getSizesAndIndexesConv2d(isNCHW, wFormat, *input, *output, bS, iC, iH, iW, oC, oH, oW, indIOioC,
indIiH, indWiC, indWmC, indWkH, indOoH);
mC = weights->sizeAt(indWmC); // channels multiplier
ConvolutionUtils::calcPadding2D(pH, pW, oH, oW, iH, iW, kH, kW, sH, sW, dH, dW, paddingMode);
std::vector<LongType> expectedWeightsShape = ConvolutionUtils::expectWeightsShape(wFormat, kH, kW, iC, mC);
REQUIRE_TRUE(weights->isSameShape(expectedWeightsShape), 0,
"DEPTHWISECONV2D CUDNN OP: wrong shape of weights array, expected is %s, but got %s instead !",
ShapeUtils::shapeAsString(expectedWeightsShape).c_str(), ShapeUtils::shapeAsString(weights).c_str());
REQUIRE_TRUE(
output->sizeAt(indIOioC) == iC * mC, 0,
"DEPTHWISECONV2D CUDNN OP: the output_channels must be equal to input_channels * channels_multiplier = %i !",
iC * mC);
if (bias)
REQUIRE_TRUE(bias->rankOf() <= 2 && oC == bias->lengthOf(), 0,
"DEPTHWISECONV2D CUDNN OP: wrong shape of array with biases, expected rank, length: <=2, %i, but got "
"%i, %i instead !",
oC, bias->rankOf(), bias->lengthOf());
std::vector<LongType> wPermut; // cudnn support format {oC, iC/groupCount, kH, kW} only, mC = 1, oC = iC (groupCount ==
// iC) that is {iC, mC, kH, kW} in our case
if (0 == wFormat)
wPermut = {2, 3, 0, 1}; // kH, kW, iC, mC -> iC, mC, kH, kW
else if (1 == wFormat)
wPermut = {1, 0, 2, 3}; // mC, iC, kH, kW -> iC, mC, kH, kW
else
wPermut = {3, 0, 1, 2}; // mC, kH, kW, iC -> iC, mC, kH, kW
std::vector<sd::LongType > perm = {iC, mC, kH, kW};
NDArray * uNewWeights = new NDArray(weights->ordering(),perm, weights->dataType(), weights->getContext());
NDArray assign = weights->permute(wPermut,false,false);
uNewWeights->assign(&assign);
std::unique_ptr<NDArray> tmpInput = {};
if (paddingMode == 1) { // in same paddingMode cudnn doesn't support asymmetric left/right top/bottopm paddings
auto ret = checkConv2dCUDNNPadAsymmetric(input, nullptr, iH, iW, oH, oW, kH, kW, sH, sW, pH, pW, dH, dW, isNCHW);
tmpInput = std::move(std::get<0>(ret));
if (tmpInput) input = tmpInput.get();
}
depthwiseConv2dCUDNN(block.launchContext(), input, uNewWeights, bias, output, kH, kW, sH, sW, pH, pW, dH, dW,
paddingMode, isNCHW);
return Status::OK;
}
//////////////////////////////////////////////////////////////////////////
PLATFORM_CHECK(depthwise_conv2d, ENGINE_CUDA) {
auto input = INPUT_VARIABLE(0); // [bS, iH, iW, iC] (NHWC) or [bS, iC, iH, iW] (NCHW)
auto weights = INPUT_VARIABLE(1); // [kH, kW, iC, mC], [mC, iC, kH, kW], [mC, kH, kW, iC]
auto bias = block.width() > 2 ? INPUT_VARIABLE(2) : nullptr; // [oC] = iC*mC
const int paddingMode = INT_ARG(8); // 0-VALID, 1-SAME, 2-CAUSAL
const int wFormat = block.getIArguments()->size() > 10
? INT_ARG(10)
: 0; // 0 - [kH, kW, iC, mC], 1 - [mC, iC, kH, kW], 2 - [mC, kH, kW, iC]
Requirements req("CUDNN DEPTHWISE_CONV2d OP");
req.expectNotEq(makeInfoVariable(paddingMode, "paddingMode"), 2) &&
req.expectEq(makeInfoVariable(weights->sizeAt(0 == wFormat ? 3 : 0), "weights#mC"), 1) &&
req.expectIn(makeInfoVariable(input->dataType(), TYPE_MSG_INPUT0),
{HALF, FLOAT32, DOUBLE}) &&
req.expectIn(makeInfoVariable(weights->dataType(), TYPE_MSG_INPUT1),
{HALF, FLOAT32, DOUBLE});
if (bias) {
req.expectIn(makeInfoVariable(bias->dataType(), TYPE_MSG_INPUT_ "#bias"),
{HALF, FLOAT32, DOUBLE});
}
req.logTheSuccess();
return req;
}
//////////////////////////////////////////////////////////////////////////
PLATFORM_IMPL(depthwise_conv2d_bp, ENGINE_CUDA) {
auto input = INPUT_VARIABLE(0); // [bS, iH, iW, iC] (NDHWC) or [bS, iC, iH, iW] (NCDHW)
auto weights = INPUT_VARIABLE(1); // [kH, kW, iC, mC], [mC, iC, kH, kW], [mC, kH, kW, iC]
auto bias = block.width() > 3 ? INPUT_VARIABLE(2) : nullptr; // [oC] = [iC*mC]
auto gradO = block.width() > 3
? INPUT_VARIABLE(3)
: INPUT_VARIABLE(2); // [bS, oH, oW, oC] (NDHWC) or [bS, oC, oH, oW] (NCDHW), epsilon_next
auto gradI = OUTPUT_VARIABLE(0); // [bS, iH, iW, iC] (NDHWC) or [bS, iC, iH, iW] (NCDHW), epsilon
auto gradW = OUTPUT_VARIABLE(1); // [kH, kW, iC, mC], [mC, iC, kH, kW], [mC, kH, kW, iC]
auto gradB = block.width() > 3 ? OUTPUT_VARIABLE(2) : nullptr; // [oC]
REQUIRE_TRUE(input->rankOf() == 4, 0,
"DEPTHWISECONV2D_BP CUDNN OP: rank of input array must be equal to 4, but got %i instead !",
input->rankOf());
REQUIRE_TRUE(weights->rankOf() == 4, 0,
"DEPTHWISECONV2D_BP CUDNN OP: rank of weights array must be equal to 4, but got %i instead !",
weights->rankOf());
REQUIRE_TRUE(gradO->rankOf() == 4, 0,
"DEPTHWISECONV2D_BP CUDNN OP: rank of output gradients (next epsilon) array must be equal to 4, but got "
"%i instead !",
gradO->rankOf());
LongType kH = INT_ARG(0) > 0 ? INT_ARG(0) : static_cast<LongType>(weights->sizeAt(0)); // filter(kernel) height
LongType kW = INT_ARG(1) > 0 ? INT_ARG(1) : static_cast<LongType>(weights->sizeAt(1)); // filter(kernel) width
LongType sH = INT_ARG(2); // strides height
LongType sW = INT_ARG(3); // strides width
LongType pH = INT_ARG(4); // paddings height
LongType pW = INT_ARG(5); // paddings width
LongType dH = INT_ARG(6); // dilations height
LongType dW = INT_ARG(7); // dilations width
int paddingMode = INT_ARG(8); // 0-VALID, 1-SAME
int isNCHW = block.getIArguments()->size() > 9 ? !INT_ARG(9) : 1; // INT_ARG(9): 1-NHWC, 0-NCHW
int wFormat = block.getIArguments()->size() > 10
? INT_ARG(10)
: 0; // 0 - [kH, kW, iC, mC], 1 - [mC, iC, kH, kW], 2 - [mC, kH, kW, iC]
LongType bS, iC, iH, iW, mC, oC, oH, oW; // batch size, input channels, input height/width, channels multiplier(oC =
// iC*mC), output channels, output height/width
LongType indIOioC, indIiH, indWmC, indWiC, indWkH, indOoH; // corresponding indexes
ConvolutionUtils::getSizesAndIndexesConv2d(isNCHW, wFormat, *input, *gradO, bS, iC, iH, iW, oC, oH, oW, indIOioC,
indIiH, indWiC, indWmC, indWkH, indOoH);
mC = weights->sizeAt(indWmC); // channels multiplier
LongType trueoH, trueoW; // correct output height, width
ConvolutionUtils::calcOutSizePool2D(trueoH, trueoW, kH, kW, sH, sW, pH, pW, dH, dW, iH, iW, paddingMode);
ConvolutionUtils::calcPadding2D(pH, pW, oH, oW, iH, iW, kH, kW, sH, sW, dH, dW, paddingMode);
std::vector<LongType> expectedGradOShape =
ShapeUtils::composeShapeUsingDimsAndIdx({bS, oC, trueoH, trueoW, 0, indIOioC, indOoH, indOoH + 1});
std::vector<LongType> expectedWeightsShape = ConvolutionUtils::expectWeightsShape(wFormat, kH, kW, iC, mC);
REQUIRE_TRUE(gradO->isSameShape(expectedGradOShape), 0,
"DEPTHWISECONV2D_BP CUDNN OP: wrong shape of output gradients (next epsilon) array, expected is %s, but "
"got %s instead !",
ShapeUtils::shapeAsString(expectedGradOShape).c_str(), ShapeUtils::shapeAsString(gradO).c_str());
REQUIRE_TRUE(weights->isSameShape(expectedWeightsShape), 0,
"DEPTHWISECONV2D_BP CUDNN OP: wrong shape of weights array, expected is %s, but got %s instead !",
ShapeUtils::shapeAsString(expectedWeightsShape).c_str(), ShapeUtils::shapeAsString(weights).c_str());
if (bias)
REQUIRE_TRUE(bias->rankOf() <= 2 && oC == bias->lengthOf(), 0,
"DEPTHWISECONV2D_BP CUDNN OP: wrong shape of array with biases, expected rank, length: <=2, %i, but "
"got %i, %i instead !",
oC, bias->rankOf(), bias->lengthOf());
std::vector<LongType> wPermut, gradWPermut; // cudnn support format {oC, iC/groupCount, kH, kW} only, mC = 1, oC = iC
// (groupCount == iC) that is {iC, mC, kH, kW}
if (0 == wFormat) {
wPermut = {2, 3, 0, 1}; // kH, kW, iC, mC -> iC, mC, kH, kW
gradWPermut = {2, 3, 0, 1}; // iC, mC, kH, kW -> kH, kW, iC, mC
} else if (1 == wFormat) {
wPermut = {1, 0, 2, 3}; // mC, iC, kH, kW -> iC, mC, kH, kW
gradWPermut = {1, 0, 2, 3}; // iC, mC, kH, kW -> mC, iC, kH, kW
} else {
wPermut = {3, 0, 1, 2}; // mC, kH, kW, iC -> iC, mC, kH, kW
gradWPermut = {1, 2, 3, 0}; // iC, mC, kH, kW -> mC, kH, kW, iC
}
std::unique_ptr<NDArray> tmpGradI = {}, tmpInput = {};
std::vector<sd::LongType> shape = {iC, mC, kH, kW};
NDArray * uNewGradW =
new NDArray(gradW->ordering(),shape, gradW->dataType(), gradW->getContext());
NDArray * uNewWeights =
new NDArray(weights->ordering(),shape, weights->dataType(), weights->getContext());
NDArray assign = weights->permute(wPermut,false,false);
uNewWeights->assign(&assign);
NDArray* newInput = input;
NDArray* newGradI = gradI;
if (paddingMode == 1) { // in same paddingMode cudnn doesn't support asymmetric left/right top/bottopm paddings
auto ret = checkConv2dCUDNNPadAsymmetric(input, gradI, iH, iW, oH, oW, kH, kW, sH, sW, pH, pW, dH, dW, isNCHW);
tmpInput = std::move(std::get<0>(ret));
tmpGradI = std::move(std::get<1>(ret));
if (tmpInput) newInput = tmpInput.get();
if (tmpGradI) newGradI = tmpGradI.get();
}
depthwiseConv2dBpCUDNN(block.launchContext(), newInput, uNewWeights, gradO, newGradI, uNewGradW, gradB,
kH, kW, sH, sW, pH, pW, dH, dW, paddingMode, isNCHW);
uNewGradW->permutei(gradWPermut,false,false);
gradW->assign(uNewGradW);
if (newInput != input) {
if (isNCHW) {
NDArray assign = (*newGradI)({0, 0, 0, 0, 0, gradI->sizeAt(2), 0, gradI->sizeAt(3)});
gradI->assign(&assign);
} else {
NDArray assign = (*newGradI)({0, 0, 0, gradI->sizeAt(1), 0, gradI->sizeAt(2), 0, 0});
gradI->assign(&assign);
}
}
return Status::OK;
}
PLATFORM_CHECK(depthwise_conv2d_bp, ENGINE_CUDA) {
auto input = INPUT_VARIABLE(0); // [bS, iH, iW, iC] (NDHWC) or [bS, iC, iH, iW] (NCDHW)
auto weights = INPUT_VARIABLE(1); // [kH, kW, iC, mC], [mC, iC, kH, kW], [mC, kH, kW, iC]
auto bias = block.width() > 3 ? INPUT_VARIABLE(2) : nullptr; // [oC] = [iC*mC]
auto gradO = block.width() > 3
? INPUT_VARIABLE(3)
: INPUT_VARIABLE(2); // [bS, oH, oW, oC] (NDHWC) or [bS, oC, oH, oW] (NCDHW), epsilon_next
const int paddingMode = INT_ARG(8); // 0-VALID, 1-SAME, 2-CAUSAL
const int isNCHW = block.getIArguments()->size() > 9 ? !INT_ARG(9) : 1; // INT_ARG(9): 0-NCHW, 1-NHWC
const int wFormat = block.getIArguments()->size() > 10
? INT_ARG(10)
: 0; // 0 - [kH, kW, iC, mC], 1 - [mC, iC, kH, kW], 2 - [mC, kH, kW, iC]
Requirements req("CUDNN DEPTHWISE_CONV2d_BP OP");
const auto inType = input->dataType();
const auto wType = weights->dataType();
const auto gType = gradO->dataType();
req.expectNotEq(makeInfoVariable(paddingMode, "paddingMode"), 2) &&
req.expectTrue(makeInfoVariable(isNCHW, "isNCHW")) &&
req.expectEq(makeInfoVariable(weights->sizeAt(0 == wFormat ? 3 : 0), "weights#mC"), 1) &&
req.expectIn(makeInfoVariable(input->dataType(), TYPE_MSG_INPUT0),
{HALF, FLOAT32, DOUBLE}) &&
req.expectIn(makeInfoVariable(weights->dataType(), TYPE_MSG_INPUT1),
{HALF, FLOAT32, DOUBLE});
if (bias) {
req.expectIn(makeInfoVariable(bias->dataType(), TYPE_MSG_INPUT_ "#bias"),
{HALF, FLOAT32, DOUBLE}) &&
req.expectIn(makeInfoVariable(gradO->dataType(), TYPE_MSG_INPUT3),
{HALF, FLOAT32, DOUBLE});
} else {
req.expectIn(makeInfoVariable(gradO->dataType(), TYPE_MSG_INPUT2),
{HALF, FLOAT32, DOUBLE});
}
req.logTheSuccess();
return req;
}
} // namespace platforms
} // namespace ops
} // namespace sd
libnd4j/include/ops/declarable/platform/cudnn/lstmLayer.cu
+673
View File
@@ -0,0 +1,673 @@
/* ******************************************************************************
*
*
* This program and the accompanying materials are made available under the
* terms of the Apache License, Version 2.0 which is available at
* https://www.apache.org/licenses/LICENSE-2.0.
*
* See the NOTICE file distributed with this work for additional
* information regarding copyright ownership.
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
* WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
* License for the specific language governing permissions and limitations
* under the License.
*
* SPDX-License-Identifier: Apache-2.0
******************************************************************************/
//
// @author AbdelRauf
//
#include <array/NDArrayFactory.h>
#include <ops/declarable/OpRegistrator.h>
#include "cudnnUtils.h"
namespace sd {
namespace ops {
namespace platforms {
// our implementation designed for 1 physical layer
constexpr int numLayers = 1;
// we will copy without using cudnnGetRNNLinLayerMatrixParams : 1 pseudo layer , isBidirectional : 2 pseudo layer
void copyWeights(const cudaStream_t &stream, bool isBidirectional, uint8_t *weightsSpace, size_t weightsSize,
uint8_t *inputWeightsData, uint8_t *recurrentWeightsData, uint8_t *biasesData, LongType inputSize,
int hiddenSize, int dataTypeSize) {
int pseudo_layer_count = isBidirectional ? 2 : 1;
uint8_t *wptr = weightsSpace;
auto wEnd = wptr + weightsSize;
// copy size for 1 full pseudo layer
// in bidirectional 1 layer consist of 2 pseduo layers
auto input_pseudo_size = 4 * inputSize * hiddenSize * dataTypeSize;
auto hidden_pseudo_size = 4 * hiddenSize * hiddenSize * dataTypeSize;
for (LongType i = 0; i < pseudo_layer_count; i++) {
if (wptr + input_pseudo_size + hidden_pseudo_size > wEnd) return;
// copy input weights
if (inputWeightsData) {
cudaMemcpyAsync(wptr, inputWeightsData, input_pseudo_size, cudaMemcpyDeviceToDevice, stream);
inputWeightsData += input_pseudo_size;
}
wptr += input_pseudo_size;
// copy recurrent weights
if (recurrentWeightsData) {
cudaMemcpyAsync(wptr, recurrentWeightsData, hidden_pseudo_size, cudaMemcpyDeviceToDevice, stream);
recurrentWeightsData += hidden_pseudo_size;
}
wptr += hidden_pseudo_size;
}
// copy bias first 4
auto bias_size = 4 * hiddenSize * dataTypeSize;
for (int i = 0; i < pseudo_layer_count; i++) {
// refill first 4 biases
if (biasesData && wptr + bias_size < wEnd) {
cudaMemcpyAsync(wptr, biasesData, bias_size, cudaMemcpyDeviceToDevice, stream);
biasesData += bias_size;
}
wptr += bias_size;
// refill next 4 with zeros
if (wptr + bias_size < wEnd) {
cudaMemsetAsync(wptr, 0, bias_size, stream);
wptr += bias_size;
}
}
// memset the rest
if (wEnd - wptr) cudaMemsetAsync(wptr, 0, wEnd - wptr, stream);
}
void cudnn_rnn_old(LaunchContext *contextPtr, int dataFormat, NDArray *input, NDArray *inputWeights,
NDArray *recurrentWeights, NDArray *biases, NDArray *prevAct, NDArray *prevMemCell,
NDArray *outputActivations, NDArray *finalTimeStepActivations, NDArray *finalMemCellState,
LongType maxSeqLength, LongType batchSize, LongType inputSize, LongType hiddenSize, double cellClip,
bool isBidirectional) {
sd_debug("cudnn rnn api %s \n", "v6");
bool training = false;
cudnnHandle_t handle = *(reinterpret_cast<cudnnHandle_t *>(contextPtr->getCuDnnHandle()));
auto stream = *(contextPtr->getCudaStream());
CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnSetStream), cudnnSetStream(handle, stream));
CudnnTensorList xDescList(maxSeqLength);
CudnnTensorList yDescList(maxSeqLength);
auto cudnnType = cudnnDataType(input->dataType());
auto dataTypeSize = input->sizeOfT();
CudnnTensor hxDesc, cxDesc, hyDesc, cyDesc;
constexpr int rankOf = 3;
const int numDirections = isBidirectional ? 2 : 1;
const int dimsX[rankOf] = {static_cast<int>(batchSize), static_cast<int>(inputSize), 1};
const int stridesX[rankOf] = {static_cast<int>(inputSize), 1, 1};
const int dimsY[rankOf] = {static_cast<int>(batchSize), static_cast<int>(hiddenSize * numDirections), 1};
const int stridesY[rankOf] = {static_cast<int>(hiddenSize * numDirections), 1, 1};
const int dimC[rankOf] = {static_cast<int>(numLayers * numDirections), static_cast<int>(batchSize), static_cast<int>(hiddenSize)};
const int strideC[rankOf] = {static_cast<int>(batchSize * hiddenSize), static_cast<int>(hiddenSize), 1};
for (int i = 0; i < maxSeqLength; i++) {
xDescList.set(i, cudnnType, rankOf, dimsX, stridesX);
yDescList.set(i, cudnnType, rankOf, dimsY, stridesY);
}
auto xDesc0 = xDescList.get(0);
hxDesc.set(cudnnType, rankOf, dimC, strideC);
cxDesc.set(cudnnType, rankOf, dimC, strideC);
hyDesc.set(cudnnType, rankOf, dimC, strideC);
cyDesc.set(cudnnType, rankOf, dimC, strideC);
PointersManager manager(contextPtr, __func__);
// dropout section
DropoutDesc dropoutDesc(nullptr);
// dropout
float dropout = 0;
size_t sizeInBytes = 0;
void *droupoutMem = nullptr;
uint64_t seed = 1; // seed
if (dropout != 0) {
dropoutDesc.create();
CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnDropoutGetStatesSize), cudnnDropoutGetStatesSize(handle, &sizeInBytes));
// allocate and set
droupoutMem = manager.allocateDevMem(sizeInBytes);
dropoutDesc.set(handle, dropout, droupoutMem, sizeInBytes, seed);
}
// RNN
RnnDesc rnnDesc;
cudnnRNNMode_t rnnCellMode = CUDNN_LSTM;
cudnnRNNAlgo_t algo = CUDNN_RNN_ALGO_STANDARD;
auto direction = isBidirectional ? CUDNN_BIDIRECTIONAL : CUDNN_UNIDIRECTIONAL;
auto mathPrec = cudnnType;
// Note: We will set some parameters manually
constexpr auto inputMode = CUDNN_LINEAR_INPUT;
rnnDesc.setUsingOldAPI(handle, inputMode, direction, rnnCellMode, algo, mathPrec, hiddenSize, numLayers, dropoutDesc);
#if CUDNN_VERSION >= CUDNN_CLIPPING_API_VER
if (cellClip > 0 && cudnnGetVersion() >= CUDNN_CLIPPING_API_VER) {
CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnRNNSetClip), cudnnRNNSetClip(handle, rnnDesc, CUDNN_RNN_CLIP_MINMAX,
CUDNN_PROPAGATE_NAN, -cellClip, cellClip));
}
#endif
// set up parameters
size_t weightsSize = 0;
CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnGetRNNParamsSize),
cudnnGetRNNParamsSize(handle, rnnDesc, xDesc0, &weightsSize, cudnnType));
FilterDesc wDesc;
int dimW[] = {static_cast<int>(weightsSize / dataTypeSize), 1, 1};
wDesc.set(cudnnType, CUDNN_TENSOR_NCHW, 3, dimW);
// allocation
void *weightsSpace = manager.allocateDevMem(weightsSize);
size_t workSpaceSizeInBytes = 0;
size_t reserveSpaceSizeInBytes = 0;
CHECK_CUDNN_FAILURE_MSG(
STRINGIZE(cudnnGetRNNWorkspaceSize),
cudnnGetRNNWorkspaceSize(handle, rnnDesc, maxSeqLength, xDescList.getDescriptors(), &workSpaceSizeInBytes));
void *workSpace = manager.allocateDevMem(workSpaceSizeInBytes);
void *reserveSpace = nullptr;
// training
if (training) {
CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnGetRNNTrainingReserveSize),
cudnnGetRNNTrainingReserveSize(handle, rnnDesc, maxSeqLength, xDescList.getDescriptors(),
&reserveSpaceSizeInBytes));
reserveSpace = manager.allocateDevMem(reserveSpaceSizeInBytes);
}
NDArray::prepareSpecialUse({outputActivations, finalTimeStepActivations, finalMemCellState},
{input, inputWeights, recurrentWeights, biases, prevAct, prevMemCell});
uint8_t *biasesData = biases ? (uint8_t *)biases->specialBuffer() : nullptr;
auto prevActData = prevAct ? prevAct->specialBuffer() : nullptr;
auto prevMemCellData = prevMemCell ? prevMemCell->specialBuffer() : nullptr;
auto finalTimeStepActivationsData = finalTimeStepActivations ? finalTimeStepActivations->specialBuffer() : nullptr;
auto finalMemCellStateData = finalMemCellState ? finalMemCellState->specialBuffer() : nullptr;
// dimension 4*nOut implies order it, ft, c't, ot
// input gate, forget gate, new gate, output gate, input gate, forget gate, new gate, output gate
// Note: our weights should be transposed and duplicated with C order to match cudnn ones
NDArray inputWeightsT, recurrentWeightsT;
uint8_t *inputWeightsData = nullptr;
uint8_t *recurrentWeightsData = nullptr;
if (inputWeights) {
inputWeightsT =
inputWeights->rankOf() == 3 ? inputWeights->permute({0, 2, 1}, 0, false).dup('c') : inputWeights->transpose().dup('c');
inputWeightsData = (uint8_t *)inputWeightsT.specialBuffer();
}
if (recurrentWeights) {
recurrentWeightsT = recurrentWeights->rankOf() == 3 ? recurrentWeights->permute({0, 2, 1}, 0, false).dup('c')
: recurrentWeights->transpose().dup('c');
recurrentWeightsData = (uint8_t *)recurrentWeightsT.specialBuffer();
}
// copy without cudnnGetRNNLinLayerMatrixParams
copyWeights(stream, isBidirectional, (uint8_t *)weightsSpace, weightsSize, inputWeightsData, recurrentWeightsData,
biasesData, inputSize, hiddenSize, dataTypeSize);
// permute based on dataformat
NDArray *argX = input;
NDArray *argOutput = outputActivations;
NDArray permutedX, outputH;
if (outputActivations != nullptr && (dataFormat != 0 || outputActivations->ordering() != 'c')) {
outputH = NDArray('c', std::vector<LongType>{maxSeqLength, batchSize, (numDirections * hiddenSize)},
outputActivations->dataType(), contextPtr);
argOutput = &outputH;
}
if (dataFormat == 1) {
permutedX = input->permute({1, 0, 2}, 0, false).dup('c');
argX = &permutedX;
}
auto xData = argX->specialBuffer();
auto yData = argOutput ? argOutput->specialBuffer() : nullptr;
if (training) {
CHECK_CUDNN_FAILURE_MSG(
STRINGIZE(cudnnRNNForwardTraining),
cudnnRNNForwardTraining(handle, rnnDesc, (int)maxSeqLength, xDescList.getDescriptors(), xData, hxDesc,
prevActData, cxDesc, prevMemCellData, wDesc, weightsSpace, yDescList.getDescriptors(),
yData, hyDesc, finalTimeStepActivationsData, cyDesc, finalMemCellStateData, workSpace,
workSpaceSizeInBytes, reserveSpace, reserveSpaceSizeInBytes));
} else {
CHECK_CUDNN_FAILURE_MSG(
STRINGIZE(cudnnRNNForwardInference),
cudnnRNNForwardInference(handle, rnnDesc, (int)maxSeqLength, xDescList.getDescriptors(), xData, hxDesc,
prevActData, cxDesc, prevMemCellData, wDesc, weightsSpace, yDescList.getDescriptors(),
yData, hyDesc, finalTimeStepActivationsData, cyDesc, finalMemCellStateData, workSpace,
workSpaceSizeInBytes));
}
// remap output
if (outputActivations != nullptr && argOutput != outputActivations) {
// refill output
if (dataFormat == 1) {
std::vector<sd::LongType> permute = {1,0,2};
NDArray assign = argOutput->permute(permute, 0, false);
outputActivations->assign(&assign);
}
}
NDArray::registerSpecialUse({outputActivations, finalTimeStepActivations, finalMemCellState},
{input, inputWeights, recurrentWeights, biases, prevAct, prevMemCell});
return;
}
#if CUDNN_VERSION >= CUDNN_NEW_RNN_API_VER
void cudnn_rnn_v8(LaunchContext *contextPtr, int dataFormat, NDArray *input, NDArray *seqLengthArray,
NDArray *inputWeights, NDArray *recurrentWeights, NDArray *biases, NDArray *prevAct,
NDArray *prevMemCell, NDArray *outputActivations, NDArray *finalTimeStepActivations,
NDArray *finalMemCellState, int maxSeqLength, int batchSize, int inputSize, int hiddenSize,
double cellClip, bool isBidirectional) {
sd_debug("cudnn rnn api %s \n", "v8");
// seqLengthArray should be int
NDArray *argSeqNdArray = nullptr;
NDArray seqArrIntData;
if (seqLengthArray) {
if (seqLengthArray->ews() == 1 && seqLengthArray->dataType() == INT32) {
argSeqNdArray = seqLengthArray;
} else {
if (seqLengthArray->dataType() != INT32) {
seqArrIntData = seqLengthArray->cast(INT32);
if (seqArrIntData.ews() != 1) seqArrIntData = seqArrIntData.dup('c');
} else {
seqArrIntData = seqLengthArray->dup('c');
}
argSeqNdArray = &seqArrIntData;
}
} else {
seqArrIntData = NDArray('c', std::vector<LongType>{batchSize}, INT32, contextPtr);
seqArrIntData.assign(maxSeqLength);
argSeqNdArray = &seqArrIntData;
}
PointersManager manager(contextPtr, __func__);
bool training = false;
cudnnHandle_t handle = *(reinterpret_cast<cudnnHandle_t *>(contextPtr->getCuDnnHandle()));
auto stream = *(contextPtr->getCudaStream());
CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnSetStream), cudnnSetStream(handle, stream));
auto cudnnType = cudnnDataType(input->dataType());
auto dataTypeSize = input->sizeOfT();
CudnnTensor hDesc, cDesc;
constexpr int rankOf = 3;
const int numDirections = isBidirectional ? 2 : 1;
const int dimC[rankOf] = {numLayers * numDirections, batchSize, hiddenSize};
const int strideC[rankOf] = {batchSize * hiddenSize, hiddenSize, 1};
hDesc.set(cudnnType, rankOf, dimC, strideC);
cDesc.set(cudnnType, rankOf, dimC, strideC);
// dropout section
DropoutDesc dropoutDesc(nullptr);
// dropout
float dropout = 0;
size_t sizeInBytes = 0;
void *droupoutMem = nullptr;
uint64_t seed = 1; // seed
if (dropout != 0) {
dropoutDesc.create();
CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnDropoutGetStatesSize), cudnnDropoutGetStatesSize(handle, &sizeInBytes));
// allocate and set
droupoutMem = manager.allocateDevMem(sizeInBytes);
dropoutDesc.set(handle, dropout, droupoutMem, sizeInBytes, seed);
}
// RNN
RnnDesc rnnDesc;
cudnnRNNMode_t rnnCellMode = CUDNN_LSTM;
cudnnRNNAlgo_t algo = CUDNN_RNN_ALGO_STANDARD;
auto direction = isBidirectional ? CUDNN_BIDIRECTIONAL : CUDNN_UNIDIRECTIONAL;
auto mathPrec = cudnnType;
// Note: We will set some parameters manually. Some of them could be parameter in future
constexpr auto inputMode = CUDNN_LINEAR_INPUT;
bool use_tensor_ops = false; // could be parameter in future
#if CUDNN_VERSION >= CUDNN_NEW_RNN_API_VER
cudnnMathType_t mathType = use_tensor_ops ? CUDNN_TENSOR_OP_MATH : CUDNN_FMA_MATH;
#else
cudnnMathType_t mathType = use_tensor_ops ? CUDNN_TENSOR_OP_MATH : CUDNN_DEFAULT_MATH;
#endif
// disable projection
int projSize = hiddenSize;
cudnnRNNBiasMode_t bias_mode = CUDNN_RNN_DOUBLE_BIAS;
uint32_t aux_flags = CUDNN_RNN_PADDED_IO_ENABLED;
rnnDesc.set(algo, rnnCellMode, bias_mode, direction, inputMode, cudnnType, mathPrec, mathType, inputSize, hiddenSize,
projSize, numLayers, dropoutDesc, aux_flags);
if (cellClip > 0) {
CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnRNNSetClip), cudnnRNNSetClip(handle, rnnDesc, CUDNN_RNN_CLIP_MINMAX,
CUDNN_PROPAGATE_NAN, -cellClip, cellClip));
}
// set Data desc
RnnDataDesc xDataDesc, yDataDesc;
bool time_major = false;
float padding_fill = 0.0f;
auto hostSeqArr = bufferInHost<int>(*argSeqNdArray);
cudnnRNNDataLayout_t layout =
dataFormat == 0 ? CUDNN_RNN_DATA_LAYOUT_SEQ_MAJOR_UNPACKED : CUDNN_RNN_DATA_LAYOUT_BATCH_MAJOR_UNPACKED;
xDataDesc.set(cudnnType, layout, maxSeqLength, batchSize, inputSize, hostSeqArr, (void *)&padding_fill);
yDataDesc.set(cudnnType, layout, maxSeqLength, batchSize, hiddenSize * numDirections, hostSeqArr,
(void *)&padding_fill);
// set up parameters
size_t weightsSize = 0;
CHECK_CUDNN_FAILURE_MSG(STRINGIZE(cudnnGetRNNWeightSpaceSize),
cudnnGetRNNWeightSpaceSize(handle, rnnDesc, &weightsSize));
// allocation
void *weightsSpace = manager.allocateDevMem(weightsSize);
// Set up work space and reserved memory
void *workSpace = nullptr;
void *reserveSpace = nullptr;
size_t workSpaceSizeInBytes = 0;
size_t reserveSpaceSizeInBytes = 0;
cudnnForwardMode_t fwdMode = training ? CUDNN_FWD_MODE_TRAINING : CUDNN_FWD_MODE_INFERENCE;
CHECK_CUDNN_FAILURE_MSG(
STRINGIZE(cudnnGetRNNTempSpaceSizes),
cudnnGetRNNTempSpaceSizes(handle, rnnDesc, fwdMode, xDataDesc, &workSpaceSizeInBytes, &reserveSpaceSizeInBytes));
workSpace = manager.allocateDevMem(workSpaceSizeInBytes);
// training
if (training) {
reserveSpace = manager.allocateDevMem(reserveSpaceSizeInBytes);
}
NDArray::prepareSpecialUse({outputActivations, finalTimeStepActivations, finalMemCellState},
{input, inputWeights, recurrentWeights, biases, prevAct, prevMemCell, argSeqNdArray});
auto xData = input->specialBuffer();
uint8_t *biasesData = biases ? (uint8_t *)biases->specialBuffer() : nullptr;
auto prevActData = prevAct ? prevAct->specialBuffer() : nullptr;
auto prevMemCellData = prevMemCell ? prevMemCell->specialBuffer() : nullptr;
auto yData = outputActivations ? outputActivations->specialBuffer() : nullptr;
auto finalTimeStepActivationsData = finalTimeStepActivations ? finalTimeStepActivations->specialBuffer() : nullptr;
auto finalMemCellStateData = finalMemCellState ? finalMemCellState->specialBuffer() : nullptr;
// dimension 4*nOut implies order it, ft, c't, ot
// input gate, forget gate, new gate, output gate, input gate, forget gate, new gate, output gate
// Note: our weights should be transposed and duplicated with C order to match cudnn ones
NDArray inputWeightsT, recurrentWeightsT;
uint8_t *inputWeightsData = nullptr;
uint8_t *recurrentWeightsData = nullptr;
if (inputWeights) {
inputWeightsT =
inputWeights->rankOf() == 3 ? inputWeights->permute({0, 2, 1}).dup('c') : inputWeights->transpose().dup('c');
inputWeightsData = (uint8_t *)inputWeightsT.specialBuffer();
}
if (recurrentWeights) {
recurrentWeightsT = recurrentWeights->rankOf() == 3 ? recurrentWeights->permute({0, 2, 1}).dup('c')
: recurrentWeights->transpose().dup('c');
recurrentWeightsData = (uint8_t *)recurrentWeightsT.specialBuffer();
}
// copy without cudnnGetRNNLinLayerMatrixParams
copyWeights(stream, isBidirectional, (uint8_t *)weightsSpace, weightsSize, inputWeightsData, recurrentWeightsData,
biasesData, inputSize, hiddenSize, dataTypeSize);
CHECK_CUDNN_FAILURE_MSG(
STRINGIZE(cudnnRNNForward),
cudnnRNNForward(handle, rnnDesc, fwdMode, (const int32_t *)argSeqNdArray->specialBuffer(), xDataDesc, xData,
yDataDesc, yData, hDesc, prevActData, finalTimeStepActivationsData, cDesc, prevMemCellData,
finalMemCellStateData, weightsSize, weightsSpace, workSpaceSizeInBytes, workSpace,
reserveSpaceSizeInBytes, reserveSpace));
NDArray::registerSpecialUse({outputActivations, finalTimeStepActivations, finalMemCellState},
{input, inputWeights, recurrentWeights, biases, prevAct, prevMemCell});
return;
}
#endif
//////////////////////////////////////////////////////////////////////////
PLATFORM_IMPL(lstmLayer, ENGINE_CUDA) {
const auto dataFormat = INT_ARG(0); // for unidirectional: 0 = [sL, bS, nIn], 1 = [bS, sL ,nIn], 2 = [bS, nIn, sL],
// for bidirectional: 3 = [sL, 2, bS, nOut] (for ONNX)
const LongType directionMode =
INT_ARG(1); // direction: 0 = fwd, 1 = bwd, 2 = bidirectional sum, 3 = bidirectional concat, 4 = bidirectional
// extra output dim (in conjunction with format dataFormat = 3)
const auto hasBiases = B_ARG(0); // indicates whether biases array is provided
const auto hasSeqLenArray = B_ARG(1); // indicates whether seqLen array is provided
const auto hasInitH = B_ARG(2); // indicates whether initial output is provided
const auto hasInitC = B_ARG(3); // indicates whether initial cell state is provided
const auto hasPH = B_ARG(4); // indicates whether peephole connections are present
const auto retFullSeq = B_ARG(5); // indicates whether to return whole time sequence h {h_0, h_1, ... , h_sL-1}
const auto retLastH = B_ARG(6); // indicates whether to return output at last time step only, in this case shape
// would be [bS, nOut] (exact shape depends on dataFormat argument)
const auto retLastC = B_ARG(7); // indicates whether to return cells state at last time step only, in this case shape
// would be [bS, nOut] (exact shape depends on dataFormat argument)
const auto cellClip = T_ARG(0); // cell clipping value, if it = 0 then do not apply clipping
const auto x = INPUT_VARIABLE(0); // input
const auto Wx = INPUT_VARIABLE(1); // input weights
const auto Wr = INPUT_VARIABLE(2); // recurrent weights
int count = 3;
const auto b = hasBiases ? INPUT_VARIABLE(count++) : nullptr; // biases
const auto seqLengthArray = hasSeqLenArray ? INPUT_VARIABLE(count++) : nullptr; // seqLen vector
const auto hI = hasInitH ? INPUT_VARIABLE(count++) : nullptr; // initial output
const auto cI = hasInitC ? INPUT_VARIABLE(count++) : nullptr; // initial cell state
const auto Wp = hasPH ? INPUT_VARIABLE(count++) : nullptr; // peephole weights
count = 0;
auto h = retFullSeq ? OUTPUT_VARIABLE(count++) : nullptr; // output
auto hL = retLastH ? OUTPUT_VARIABLE(count++) : nullptr; // output at last step
auto cL = retLastC ? OUTPUT_VARIABLE(count++) : nullptr; // cell state at last step
REQUIRE_TRUE(cellClip >= 0, 0, "LSTM_LAYER operation: cell clipping value should be nonnegative (>=0) !");
REQUIRE_TRUE(retFullSeq || retLastH || retLastC, 0,
"LSTM_LAYER operation: please specify what output arrays to produce !");
// evaluate dimensions
const LongType seqLength = dataFormat == 3 ? x->sizeAt(0) : x->sizeAt(dataFormat);
const LongType bS = dataFormat == 1 || dataFormat == 2 ? x->sizeAt(0) : x->sizeAt(1);
const LongType nIn = dataFormat == 2 ? x->sizeAt(1) : x->sizeAt(2);
const LongType nOut = Wx->sizeAt(-1) / 4;
const LongType hiddenSize = nOut;
auto contextPtr = block.launchContext();
bool isBidirectional = directionMode >= 2;
if (!isBidirectional) { // no bidirectional
// Wx validation
if (Wx->rankOf() != 2 || Wx->sizeAt(0) != nIn)
REQUIRE_TRUE(false, 0, "LSTM_LAYER operation: wrong shape of input weights, expected is %s, but got %s instead !",
ShapeUtils::shapeAsString({nIn, 4 * nOut}).c_str(), ShapeUtils::shapeAsString(Wx).c_str());
// Wr validation
if (Wr->rankOf() != 2 || Wr->sizeAt(0) != nOut || Wr->sizeAt(1) != 4 * nOut)
REQUIRE_TRUE(false, 0,
"LSTM_LAYER operation: wrong shape of recurrent weights, expected is %s, but got %s instead !",
ShapeUtils::shapeAsString({nOut, 4 * nOut}).c_str(), ShapeUtils::shapeAsString(Wr).c_str());
// biases validation
if (b != nullptr && (b->rankOf() != 1 || b->sizeAt(0) != 4 * nOut))
REQUIRE_TRUE(false, 0, "LSTM_LAYER operation: wrong shape of biases, expected is %s, but got %s instead !",
ShapeUtils::shapeAsString({4 * nOut}).c_str(), ShapeUtils::shapeAsString(b).c_str());
// initial output validation
if (hI != nullptr && (hI->rankOf() != 2 || hI->sizeAt(0) != bS || hI->sizeAt(1) != nOut))
REQUIRE_TRUE(false, 0,
"LSTM_LAYER operation: wrong shape of initial output, expected is %s, but got %s instead !",
ShapeUtils::shapeAsString({bS, nOut}).c_str(), ShapeUtils::shapeAsString(hI).c_str());
// initial cell validation
if (cI != nullptr && (cI->rankOf() != 2 || cI->sizeAt(0) != bS || cI->sizeAt(1) != nOut))
REQUIRE_TRUE(false, 0,
"LSTM_LAYER operation: wrong shape of initial cell state, expected is %s, but got %s instead !",
ShapeUtils::shapeAsString({bS, nOut}).c_str(), ShapeUtils::shapeAsString(cI).c_str());
} else { // bidirectional
// Wx validation
if (Wx->rankOf() != 3 || Wx->sizeAt(0) != 2 || Wx->sizeAt(1) != nIn)
REQUIRE_TRUE(false, 0, "LSTM_LAYER operation: wrong shape of input weights, expected is %s, but got %s instead !",
ShapeUtils::shapeAsString({2, nIn, 4 * nOut}).c_str(), ShapeUtils::shapeAsString(Wx).c_str());
// Wr validation
if (Wr->rankOf() != 3 || Wr->sizeAt(0) != 2 || Wr->sizeAt(1) != nOut || Wr->sizeAt(2) != 4 * nOut)
REQUIRE_TRUE(false, 0,
"LSTM_LAYER operation: wrong shape of recurrent weights, expected is %s, but got %s instead !",
ShapeUtils::shapeAsString({2, nOut, 4 * nOut}).c_str(), ShapeUtils::shapeAsString(Wr).c_str());
// biases validation
if (b != nullptr && (b->rankOf() != 2 || b->sizeAt(0) != 2 || b->sizeAt(1) != 4 * nOut))
REQUIRE_TRUE(false, 0, "LSTM_LAYER operation: wrong shape of biases, expected is %s, but got %s instead !",
ShapeUtils::shapeAsString({2, 4 * nOut}).c_str(), ShapeUtils::shapeAsString(b).c_str());
// initial output validation
if (hI != nullptr && (hI->rankOf() != 3 || hI->sizeAt(0) != 2 || hI->sizeAt(1) != bS || hI->sizeAt(2) != nOut))
REQUIRE_TRUE(false, 0,
"LSTM_LAYER operation: wrong shape of initial output, expected is %s, but got %s instead !",
ShapeUtils::shapeAsString({2, bS, nOut}).c_str(), ShapeUtils::shapeAsString(hI).c_str());
// initial cell validation
if (cI != nullptr && (cI->rankOf() != 3 || cI->sizeAt(0) != 2 || cI->sizeAt(1) != bS || cI->sizeAt(2) != nOut))
REQUIRE_TRUE(false, 0,
"LSTM_LAYER operation: wrong shape of initial cell state, expected is %s, but got %s instead !",
ShapeUtils::shapeAsString({2, bS, nOut}).c_str(), ShapeUtils::shapeAsString(cI).c_str());
}
#if CUDNN_VERSION < CUDNN_NEW_RNN_API_VER
cudnn_rnn_old(contextPtr, dataFormat, x, Wx, Wr, b, hI, cI, h, hL, cL, seqLength, bS, nIn, hiddenSize,
(double)cellClip, isBidirectional);
#else
if (cudnnGetVersion() >= CUDNN_NEW_RNN_API_VER) {
cudnn_rnn_v8(contextPtr, dataFormat, x, seqLengthArray, Wx, Wr, b, hI, cI, h, hL, cL, seqLength, bS, nIn,
hiddenSize, (double)cellClip, isBidirectional);
} else {
cudnn_rnn_old(contextPtr, dataFormat, x, Wx, Wr, b, hI, cI, h, hL, cL, seqLength, bS, nIn, hiddenSize,
(double)cellClip, isBidirectional);
}
#endif
return Status::OK;
}
// Cudnn Lstm:
// Forward inference implemented using v6, and v8 (when version > 8.0.1) api calls.
// As our Cuda Lstm implementation has 1 layer. Cudnn implementation was implemented for 1 physical layer
// Cudnn helper restrictions:
// - all NDArrays should be the same type
// - dataFormat should be 0 or 1
// - only unidirectional (directionMode == 0) and bidirectional concat (directionMode == 3)
// - no peephole connection
// - Clipping is allowed for cudnn version >= 7.2.1
// - SeqLen array is allowed for cudnn version >= 8.0.1
// - gateActivation: sigmoid, cellActivation and outputActivation: tanh
// - NDArrays (excluding the weight arrays, as we have to transpose or permute it) should follow 'c' order and ews()==1
PLATFORM_CHECK(lstmLayer, ENGINE_CUDA) {
const auto dataFormat = INT_ARG(0); // for unidirectional: 0 = [sL, bS, nIn], 1 = [bS, sL ,nIn], 2 = [bS, nIn, sL],
// for bidirectional: 3 = [sL, 2, bS, nOut] (for ONNX)
const auto directionMode =
INT_ARG(1); // direction: 0 = fwd, 1 = bwd, 2 = bidirectional sum, 3 = bidirectional concat, 4 = bidirectional
// extra output dim (in conjunction with format dataFormat = 3)
// integer numbers corresponding to activations: 0=tanh, 1=relu, 2=sigmoid, 3=affine, 4=leaky relu, 5= thresholded
// relu, 6=scaled tanh, 7=hard sigmoid, 8=ELU, 9=softsign, 10=softplus
const auto gateAct = INT_ARG(2); // activation for input (i), forget (f) and output (o) gates
const auto cellAct = INT_ARG(3); // activation for cell state (c)
const auto outAct = INT_ARG(4); // activation for output (h)
const auto hasBiases = B_ARG(0); // indicates whether biases array is provided
const auto hasSeqLenArray = B_ARG(1); // indicates whether seqLen array is provided
const auto hasInitH = B_ARG(2); // indicates whether initial output is provided
const auto hasInitC = B_ARG(3); // indicates whether initial cell state is provided
const auto hasPH = B_ARG(4); // indicates whether peephole connections are present
const auto retFullSeq = B_ARG(5); // indicates whether to return whole time sequence h {h_0, h_1, ... , h_sL-1}
const auto retLastH = B_ARG(6); // indicates whether to return output at last time step only, in this case shape
// would be [bS, nOut] (exact shape depends on dataFormat argument)
const auto retLastC = B_ARG(7); // indicates whether to return cells state at last time step only, in this case shape
// would be [bS, nOut] (exact shape depends on dataFormat argument)
const auto cellClip = T_ARG(0); // cell clipping value, if it = 0 then do not apply clipping
const auto x = INPUT_VARIABLE(0); // input
const auto Wx = INPUT_VARIABLE(1); // input weights
const auto Wr = INPUT_VARIABLE(2); // recurrent weights
int count = 3;
const auto b = hasBiases ? INPUT_VARIABLE(count++) : nullptr; // biases
const auto hI = hasInitH ? INPUT_VARIABLE(count++) : nullptr; // initial output
const auto cI = hasInitC ? INPUT_VARIABLE(count++) : nullptr; // initial cell state
count = 0;
auto h = retFullSeq ? OUTPUT_VARIABLE(count++) : nullptr; // output
auto hL = retLastH ? OUTPUT_VARIABLE(count++) : nullptr; // output at last step
auto cL = retLastC ? OUTPUT_VARIABLE(count++) : nullptr; // cell state at last step
DataType xType = x->dataType();
DataType WxType = Wx->dataType();
DataType WrType = Wr->dataType();
Requirements req("CUDNN LSTMLAYER OP");
// cudnn related restrictions //gateAct: sigmoid, cellAct: tanh adn et cetera
// integer numbers corresponding to activations: 0=tanh, 1=relu, 2=sigmoid, 3=affine,
// 4=leaky relu, 5= thresholded relu, 6=scaled tanh, 7=hard sigmoid, 8=ELU, 9=softsign, 10=softplus
req.expectEq(makeInfoVariable(gateAct, "gate Activation"), makeInfoVariable(2, "sigmoid")) &&
req.expectEq(makeInfoVariable(cellAct, "cell Activation"), makeInfoVariable(2, "tanh")) &&
req.expectEq(makeInfoVariable(outAct, "out Activation"), makeInfoVariable(2, "tanh")) &&
req.expectFalse(makeInfoVariable(hasPH, HAVE_PEEPHOLE), EXPECTED_NOT_SUPPORTED) &&
req.expectIn(makeInfoVariable(directionMode, "directionMode"), {0, 3}) &&
req.expectIn(makeInfoVariable(dataFormat, "data Format"), {0, 1});
if (req) {
// cudnn api version related restrictions in our helpers
size_t cudnn_version = cudnnGetVersion();
// though seqlengthArray was added in earlier versions we do not handle it below 8.0.0.1
#if CUDNN_VERSION < CUDNN_NEW_RNN_API_VER
// implRestrictions = implRestrictions && !hasSeqLenArray;
req.expectFalse(makeInfoVariable(hasSeqLenArray, HAVE_SEQLENARR), EXPECTED_NOT_SUPPORTED);
#else
// implRestrictions = implRestrictions && (cudnn_version >= CUDNN_NEW_RNN_API_VER || !hasSeqLenArray);
if (cudnn_version < CUDNN_NEW_RNN_API_VER) {
req.expectFalse(makeInfoVariable(hasSeqLenArray, HAVE_SEQLENARR), EXPECTED_NOT_SUPPORTED);
}
#endif
// implRestrictions = implRestrictions && (cudnn_version >= CUDNN_CLIPPING_API_VER || cellClip==0);
if (cudnn_version < CUDNN_CLIPPING_API_VER) {
req.expectEq(makeInfoVariable(cellClip, MSG_CELL_CLIPPING), 0);
}
}
// restriction that comes either from not setting Descriptor or not handling manipulation:
// restrict0: the same types
req.expectEq(makeInfoVariable(x->ordering(), ORDERING_MSG_INPUT0), 'c') &&
req.expectEq(makeInfoVariable(WxType, TYPE_MSG_INPUT1), makeInfoVariable(xType, TYPE_MSG_INPUT0)) &&
req.expectEq(makeInfoVariable(WrType, TYPE_MSG_INPUT2), makeInfoVariable(xType, TYPE_MSG_INPUT0));
if (b)
req.expectEq(makeInfoVariable(b->dataType(), TYPE_MSG_INPUT_ "#bias"), makeInfoVariable(xType, TYPE_MSG_INPUT0));
if (hI) {
req.expectEq(makeInfoVariable(hI->dataType(), TYPE_MSG_INPUT_ "#hI"), makeInfoVariable(xType, TYPE_MSG_INPUT0)) &&
req.expectEq(makeInfoVariable(hI->ordering(), ORDERING_MSG_INPUT_ "#hI"), 'c') &&
}
if (cI) {
req.expectEq(makeInfoVariable(cI->dataType(), TYPE_MSG_INPUT_ "#cI"), makeInfoVariable(xType, TYPE_MSG_INPUT0)) &&
req.expectEq(makeInfoVariable(cI->ordering(), ORDERING_MSG_INPUT_ "#cI"), 'c') &&
}
if (h) {
req.expectEq(makeInfoVariable(h->dataType(), TYPE_MSG_OUTPUT_ "#h"), makeInfoVariable(xType, TYPE_MSG_INPUT0)) &&
req.expectEq(makeInfoVariable(h->ordering(), ORDERING_MSG_OUTPUT_ "#h"), 'c') &&
}
if (hL) {
req.expectEq(makeInfoVariable(hL->dataType(), TYPE_MSG_OUTPUT_ "#hL"), makeInfoVariable(xType, TYPE_MSG_INPUT0)) &&
req.expectEq(makeInfoVariable(hL->ordering(), ORDERING_MSG_OUTPUT_ "#hL"), 'c') &&
}
if (cL) {
req.expectEq(makeInfoVariable(cL->dataType(), TYPE_MSG_OUTPUT_ "#cL"), makeInfoVariable(xType, TYPE_MSG_INPUT0)) &&
req.expectEq(makeInfoVariable(cL->ordering(), ORDERING_MSG_OUTPUT_ "#cL"), 'c') &&
}
req.logTheSuccess();
return req;
}
} // namespace platforms
} // namespace ops
} // namespace sd
libnd4j/include/ops/declarable/platform/cudnn/maxpool2d.cu
+157
View File
@@ -0,0 +1,157 @@
/* ******************************************************************************
*
*
* This program and the accompanying materials are made available under the
* terms of the Apache License, Version 2.0 which is available at
* https://www.apache.org/licenses/LICENSE-2.0.
*
* See the NOTICE file distributed with this work for additional
* information regarding copyright ownership.
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
* WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
* License for the specific language governing permissions and limitations
* under the License.
*
* SPDX-License-Identifier: Apache-2.0
******************************************************************************/
//
// @author Yurii Shyrma (iuriish@yahoo.com)
//
#include <ops/declarable/helpers/convolutions.h>
#include "cudnnUtils.h"
namespace sd {
namespace ops {
namespace platforms {
//////////////////////////////////////////////////////////////////////////
PLATFORM_IMPL(maxpool2d, ENGINE_CUDA) {
auto input = INPUT_VARIABLE(0);
auto output = OUTPUT_VARIABLE(0);
// 0,1 - kernel Height/Width; 2,3 - stride Height/Width; 4,5 - pad Height/Width; 6,7 - dilation Height/Width; 8 -
// paddingModee;
const LongType kH = INT_ARG(0);
const LongType kW = INT_ARG(1);
const LongType sH = INT_ARG(2);
const LongType sW = INT_ARG(3);
LongType pH = INT_ARG(4);
LongType pW = INT_ARG(5);
const LongType dH = INT_ARG(6);
const LongType dW = INT_ARG(7);
const auto paddingMode = static_cast<bool>(INT_ARG(8));
const int isNCHW = block.getIArguments()->size() > 10 ? !INT_ARG(10) : 1; // INT_ARG(10): 0-NCHW, 1-NHWC
REQUIRE_TRUE(input->rankOf() == 4, 0, "MAXPOOL2D CUDNN op: input should have rank of 4, but got %i instead",
input->rankOf());
REQUIRE_TRUE(dH != 0 && dW != 0, 0, "MAXPOOL2D CUDNN op: dilation must not be zero, but got instead {%i, %i}", dH,
dW);
LongType oH = 0;
LongType oW = 0;
const LongType iH = static_cast<LongType>(isNCHW ? input->sizeAt(2) : input->sizeAt(1));
const LongType iW = static_cast<LongType>(isNCHW ? input->sizeAt(3) : input->sizeAt(2));
ConvolutionUtils::calcOutSizePool2D(oH, oW, kH, kW, sH, sW, pH, pW, dH, dW, iH, iW, paddingMode);
if (paddingMode) ConvolutionUtils::calcPadding2D(pH, pW, oH, oW, iH, iW, kH, kW, sH, sW, dH, dW);
pooling2dCUDNN(block.launchContext(), input, output, kH, kW, sH, sW, pH, pW, dH, dW, isNCHW, CUDNN_POOLING_MAX);
return Status::OK;
}
//////////////////////////////////////////////////////////////////////////
PLATFORM_CHECK(maxpool2d, ENGINE_CUDA) {
auto input = INPUT_VARIABLE(0);
auto output = OUTPUT_VARIABLE(0);
Requirements req("CUDNN MAXPOOL2d OP");
req.expectEq(makeInfoVariable(input->dataType(), TYPE_MSG_INPUT),
makeInfoVariable(output->dataType(), TYPE_MSG_OUTPUT)) &&
req.expectIn(makeInfoVariable(input->dataType(), TYPE_MSG_INPUT),
{INT32, HALF, FLOAT32, DOUBLE});
req.logTheSuccess();
return req;
}
//////////////////////////////////////////////////////////////////////////
PLATFORM_IMPL(maxpool2d_bp, ENGINE_CUDA) {
auto input = INPUT_VARIABLE(0); // [bS, iH, iW, iC] (NHWC) or [bS, iC, iH, iW] (NCHW)
auto gradO = INPUT_VARIABLE(1); // [bS, oH, oW, oC] (NHWC) or [bS, oC, oH, oW] (NCHW), epsilon_next
auto gradI = OUTPUT_VARIABLE(0); // [bS, iH, iW, iC] (NHWC) or [bS, iC, iH, iW] (NCHW), epsilon
const LongType kH = INT_ARG(0); // filter(kernel) height
const LongType kW = INT_ARG(1); // filter(kernel) width
const LongType sH = INT_ARG(2); // strides height
const LongType sW = INT_ARG(3); // strides width
LongType pH = INT_ARG(4); // paddings height
LongType pW = INT_ARG(5); // paddings width
const LongType dH = INT_ARG(6); // dilations height
const LongType dW = INT_ARG(7); // dilations width
const auto paddingMode = INT_ARG(8); // 0-VALID, 1-SAME
const auto isNCHW = block.getIArguments()->size() > 10 ? !INT_ARG(10) : 1; // INT_ARG(10): 0-NCHW, 1-NHWC
REQUIRE_TRUE(input->rankOf() == 4, 0, "MAXPOOL2D_BP CUDNN op: input should have rank of 4, but got %i instead",
input->rankOf());
REQUIRE_TRUE(dH != 0 && dW != 0, 0, "MAXPOOL2D_BP CUDNN op: dilation must not be zero, but got instead {%i, %i}", dH,
dW);
LongType bS, iC, iH, iW, oC, oH,
oW; // batch size, input channels, input height/width, output channels, output height/width;
LongType indIOioC, indIiH, indWoC, indWiC, indWkH, indOoH; // corresponding indexes
ConvolutionUtils::getSizesAndIndexesConv2d(isNCHW, 0, *input, *gradO, bS, iC, iH, iW, oC, oH, oW, indIOioC, indIiH,
indWiC, indWoC, indWkH, indOoH);
std::vector<LongType> expectedGradOShape =
ShapeUtils::composeShapeUsingDimsAndIdx({bS, iC, oH, oW, 0, indIOioC, indIiH, indIiH + 1});
std::vector<LongType> expectedGradIShape =
ShapeUtils::composeShapeUsingDimsAndIdx({bS, iC, iH, iW, 0, indIOioC, indIiH, indIiH + 1});
REQUIRE_TRUE(gradO->isSameShape(expectedGradOShape), 0,
"MAXPOOL2D_BP CUDNN op: wrong shape of output's gradients array (next epsilon), expected is %s, but got "
"%s instead !",
ShapeUtils::shapeAsString(expectedGradOShape).c_str(), ShapeUtils::shapeAsString(gradO).c_str());
REQUIRE_TRUE(
gradI->isSameShape(expectedGradIShape), 0,
"MAXPOOL2D_BP CUDNN op: wrong shape of input's gradients array (epsilon), expected is %s, but got %s instead !",
ShapeUtils::shapeAsString(expectedGradIShape).c_str(), ShapeUtils::shapeAsString(gradI).c_str());
if (paddingMode) // SAME
ConvolutionUtils::calcPadding2D(pH, pW, oH, oW, iH, iW, kH, kW, sH, sW, dH, dW);
pooling2dBpCUDNN(block.launchContext(), input, gradO, gradI, kH, kW, sH, sW, pH, pW, dH, dW, isNCHW,
CUDNN_POOLING_MAX);
return Status::OK;
}
PLATFORM_CHECK(maxpool2d_bp, ENGINE_CUDA) {
auto input = INPUT_VARIABLE(0); // [bS, iH, iW, iC] (NHWC) or [bS, iC, iH, iW] (NCHW)
auto gradO = INPUT_VARIABLE(1); // [bS, oH, oW, oC] (NHWC) or [bS, oC, oH, oW] (NCHW), epsilon_next
auto gradI = OUTPUT_VARIABLE(0); // [bS, iH, iW, iC] (NHWC) or [bS, iC, iH, iW] (NCHW), epsilon
Requirements req("CUDNN MAXPOOL2d_BP OP");
req.expectEq(makeInfoVariable(input->ordering(), ORDERING_MSG_INPUT), 'c') &&
req.expectEq(makeInfoVariable(input->dataType(), TYPE_MSG_INPUT0),
makeInfoVariable(gradO->dataType(), TYPE_MSG_INPUT1)) &&
req.expectEq(makeInfoVariable(input->dataType(), TYPE_MSG_INPUT),
makeInfoVariable(gradI->dataType(), TYPE_MSG_OUTPUT)) &&
req.expectIn(makeInfoVariable(input->dataType(), TYPE_MSG_INPUT),
{INT32, HALF, FLOAT32, DOUBLE}) &&
req.expect(
makeShapeInfoVariable(input, SHAPE_MSG_INPUT0), makeShapeInfoVariable(gradI, SHAPE_MSG_OUTPUT),
[](const decltype(input)& l, const decltype(gradI)& r) {
return shape::haveSameShapeAndStrides(l->shapeInfo(), r->shapeInfo());
},
EXPECTED_EQ_MSG);
req.logTheSuccess();
return req;
}
} // namespace platforms
} // namespace ops
} // namespace sd
libnd4j/include/ops/declarable/platform/cudnn/maxpool3d.cu
+171
View File
@@ -0,0 +1,171 @@
/* ******************************************************************************
*
*
* This program and the accompanying materials are made available under the
* terms of the Apache License, Version 2.0 which is available at
* https://www.apache.org/licenses/LICENSE-2.0.
*
* See the NOTICE file distributed with this work for additional
* information regarding copyright ownership.
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
* WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
* License for the specific language governing permissions and limitations
* under the License.
*
* SPDX-License-Identifier: Apache-2.0
******************************************************************************/
//
// @author Yurii Shyrma (iuriish@yahoo.com)
//
#include <ops/declarable/helpers/convolutions.h>
#include "cudnnUtils.h"
namespace sd {
namespace ops {
namespace platforms {
//////////////////////////////////////////////////////////////////////////
PLATFORM_IMPL(maxpool3dnew, ENGINE_CUDA) {
auto input = INPUT_VARIABLE(0); // [bS, iD, iH, iW, iC] (NDHWC) or [bS, iC, iD, iH, iW] (NCDHW)
auto output = OUTPUT_VARIABLE(0); // [bS, oD, oH, oW, iC] (NDHWC) or [bS, iC, oD, oH, oW] (NCDHW)
LongType kD = INT_ARG(0); // filter(kernel) depth
LongType kH = INT_ARG(1); // filter(kernel) height
LongType kW = INT_ARG(2); // filter(kernel) width
LongType sD = INT_ARG(3); // strides depth
LongType sH = INT_ARG(4); // strides height
LongType sW = INT_ARG(5); // strides width
LongType pD = INT_ARG(6); // paddings depth
LongType pH = INT_ARG(7); // paddings height
LongType pW = INT_ARG(8); // paddings width
LongType dD = INT_ARG(9); // dilations depth
LongType dH = INT_ARG(10); // dilations height
LongType dW = INT_ARG(11); // dilations width
int paddingMode = INT_ARG(12); // 1-SAME, 0-VALID
// int extraParam0 = INT_ARG(13);
int isNCDHW = block.getIArguments()->size() > 14 ? !INT_ARG(14) : 1; // 0-NCDHW, 1-NDHWC
REQUIRE_TRUE(input->rankOf() == 5, 0,
"MAXPOOL3DNEW CUDNN OP: rank of input array must be equal to 5, but got %i instead !", input->rankOf());
REQUIRE_TRUE(dD != 0 && dH != 0 && dW != 0, 0,
"MAXPOOL3DNEW CUDNN OP: dilation must not be zero, but got instead {%i, %i, %i}", dD, dH, dW);
LongType bS, iC, iD, iH, iW, oC, oD, oH,
oW; // batch size, input channels, input depth/height/width, output channels, output depth/height/width;
LongType indIOioC, indIOioD, indWoC, indWiC, indWkD; // corresponding indexes
ConvolutionUtils::getSizesAndIndexesConv3d(isNCDHW, 0, *input, *output, bS, iC, iD, iH, iW, oC, oD, oH, oW, indIOioC,
indIOioD, indWiC, indWoC, indWkD);
std::vector<LongType> expectedOutputShape =
ShapeUtils::composeShapeUsingDimsAndIdx({bS, iC, oD, oH, oW, 0, indIOioC, indIOioD, indIOioD + 1, indIOioD + 2});
REQUIRE_TRUE(output->isSameShape(expectedOutputShape), 0,
"MAXPOOL3DNEW CUDNN OP: wrong shape of output array, expected is %s, but got %s instead !",
ShapeUtils::shapeAsString(expectedOutputShape).c_str(), ShapeUtils::shapeAsString(output).c_str());
if (paddingMode) // SAME
ConvolutionUtils::calcPadding3D(pD, pH, pW, oD, oH, oW, iD, iH, iW, kD, kH, kW, sD, sH, sW, dD, dH, dW);
pooling3dCUDNN(block.launchContext(), input, output, kD, kH, kW, sD, sH, sW, pD, pH, pW, dD, dH, dW, isNCDHW,
CUDNN_POOLING_MAX);
return Status::OK;
}
//////////////////////////////////////////////////////////////////////////
PLATFORM_CHECK(maxpool3dnew, ENGINE_CUDA) {
auto input = INPUT_VARIABLE(0);
auto output = OUTPUT_VARIABLE(0);
Requirements req("CUDNN MAXPOOL3d OP");
req.expectEq(makeInfoVariable(input->dataType(), TYPE_MSG_INPUT),
makeInfoVariable(output->dataType(), TYPE_MSG_OUTPUT)) &&
req.expectIn(makeInfoVariable(input->dataType(), TYPE_MSG_INPUT),
{INT32, HALF, FLOAT32, DOUBLE});
req.logTheSuccess();
return req;
}
//////////////////////////////////////////////////////////////////////////
PLATFORM_IMPL(maxpool3dnew_bp, ENGINE_CUDA) {
auto input = INPUT_VARIABLE(0); // [bS, iD, iH, iW, iC] (NDHWC) or [bS, iC, iD, iH, iW] (NCDHW)
auto gradO = INPUT_VARIABLE(1); // [bS, oD, oH, oW, oC] (NDHWC) or [bS, oC, oD, oH, oW] (NCDHW), epsilon_next
auto gradI = OUTPUT_VARIABLE(0); // [bS, iD, iH, iW, iC] (NDHWC) or [bS, iC, iD, iH, iW] (NCDHW), epsilon
const LongType kD = INT_ARG(0); // filter(kernel) depth
const LongType kH = INT_ARG(1); // filter(kernel) height
const LongType kW = INT_ARG(2); // filter(kernel) width
const LongType sD = INT_ARG(3); // strides depth
const LongType sH = INT_ARG(4); // strides height
const LongType sW = INT_ARG(5); // strides width
LongType pD = INT_ARG(6); // paddings depth
LongType pH = INT_ARG(7); // paddings height
LongType pW = INT_ARG(8); // paddings width
const LongType dD = INT_ARG(9); // dilations depth
const LongType dH = INT_ARG(10); // dilations height
const LongType dW = INT_ARG(11); // dilations width
const int isSameMode = INT_ARG(12); // 1-SAME, 0-VALID
// const int extraParam0 = INT_ARG(13); // define what divisor to use while
// averaging
const int isNCDHW = block.getIArguments()->size() > 14 ? !INT_ARG(14) : 1; // 0-NCDHW, 1-NDHWC
REQUIRE_TRUE(input->rankOf() == 5, 0, "MAXPOOL3DNEW_BP CUDNN OP: input should have rank of 5, but got %i instead",
input->rankOf());
REQUIRE_TRUE(dD != 0 && dH != 0 && dW != 0, 0,
"MAXPOOL3DNEW_BP CUDNN OP: dilation must not be zero, but got instead {%i, %i, %i}", dD, dH, dW);
LongType bS, iC, iD, iH, iW, oC, oD, oH,
oW; // batch size, input channels, input depth/height/width, output channels, output depth/height/width;
LongType indIOioC, indIOioD, indWoC, indWiC, indWkD; // corresponding indexes
ConvolutionUtils::getSizesAndIndexesConv3d(isNCDHW, 0, *input, *gradO, bS, iC, iD, iH, iW, oC, oD, oH, oW, indIOioC,
indIOioD, indWiC, indWoC, indWkD);
std::vector<LongType> expectedGradOShape =
ShapeUtils::composeShapeUsingDimsAndIdx({bS, iC, oD, oH, oW, 0, indIOioC, indIOioD, indIOioD + 1, indIOioD + 2});
std::vector<LongType> expectedGradIShape =
ShapeUtils::composeShapeUsingDimsAndIdx({bS, iC, iD, iH, iW, 0, indIOioC, indIOioD, indIOioD + 1, indIOioD + 2});
REQUIRE_TRUE(gradO->isSameShape(expectedGradOShape), 0,
"MAXPOOL3DNEW_BP CUDNN: wrong shape of output's gradients array (next epsilon), expected is %s, but got "
"%s instead !",
ShapeUtils::shapeAsString(expectedGradOShape).c_str(), ShapeUtils::shapeAsString(gradO).c_str());
REQUIRE_TRUE(
gradI->isSameShape(expectedGradIShape), 0,
"MAXPOOL3DNEW_BP CUDNN: wrong shape of input's gradients array (epsilon), expected is %s, but got %s instead !",
ShapeUtils::shapeAsString(expectedGradIShape).c_str(), ShapeUtils::shapeAsString(gradI).c_str());
if (isSameMode) // SAME
ConvolutionUtils::calcPadding3D(pD, pH, pW, oD, oH, oW, iD, iH, iW, kD, kH, kW, sD, sH, sW, dD, dH, dW);
pooling3dBpCUDNN(block.launchContext(), input, gradO, gradI, kD, kH, kW, sD, sH, sW, pD, pH, pW, dD, dH, dW, isNCDHW,
CUDNN_POOLING_MAX);
return Status::OK;
}
PLATFORM_CHECK(maxpool3dnew_bp, ENGINE_CUDA) {
auto input = INPUT_VARIABLE(0); // [bS, iD, iH, iW, iC] (NDHWC) or [bS, iC, iD, iH, iW] (NCDHW)
auto gradO = INPUT_VARIABLE(1); // [bS, oD, oH, oW, oC] (NDHWC) or [bS, oC, oD, oH, oW] (NCDHW), epsilon_next
auto gradI = OUTPUT_VARIABLE(0); // [bS, iD, iH, iW, iC] (NDHWC) or [bS, iC, iD, iH, iW] (NCDHW), epsilon
Requirements req("CUDNN MAXPOOL3d_BP OP");
req.expectEq(makeInfoVariable(input->dataType(), TYPE_MSG_INPUT0),
makeInfoVariable(gradO->dataType(), TYPE_MSG_INPUT1)) &&
req.expectEq(makeInfoVariable(input->dataType(), TYPE_MSG_INPUT),
makeInfoVariable(gradI->dataType(), TYPE_MSG_OUTPUT)) &&
req.expectIn(makeInfoVariable(input->dataType(), TYPE_MSG_INPUT),
{INT32, HALF, FLOAT32, DOUBLE}) &&
req.expect(
makeShapeInfoVariable(input, SHAPE_MSG_INPUT0), makeShapeInfoVariable(gradI, SHAPE_MSG_OUTPUT),
[](const decltype(input)& l, const decltype(gradI)& r) {
return shape::haveSameShapeAndStrides(l->shapeInfo(), r->shapeInfo());
},
EXPECTED_EQ_MSG);
req.logTheSuccess();
return req;
}
} // namespace platforms
} // namespace ops
} // namespace sd
libnd4j/include/ops/declarable/platform/mkldnn/avgpooling2d.cpp
+152
View File
@@ -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
libnd4j/include/ops/declarable/platform/mkldnn/avgpooling3d.cpp
+157
View File
@@ -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
libnd4j/include/ops/declarable/platform/mkldnn/batchnorm.cpp
+639
View File
@@ -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
libnd4j/include/ops/declarable/platform/mkldnn/concat.cpp
+197
View File
@@ -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
libnd4j/include/ops/declarable/platform/mkldnn/conv2d.cpp
+439
View File
@@ -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
libnd4j/include/ops/declarable/platform/mkldnn/conv3d.cpp
+458
View File
@@ -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
libnd4j/include/ops/declarable/platform/mkldnn/deconv2d.cpp
+532
View File
@@ -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
libnd4j/include/ops/declarable/platform/mkldnn/deconv2d_tf.cpp
+232
View File
@@ -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
libnd4j/include/ops/declarable/platform/mkldnn/deconv3d.cpp
+546
View File
@@ -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
libnd4j/include/ops/declarable/platform/mkldnn/depthwiseConv2d.cpp
+570
View File
@@ -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
libnd4j/include/ops/declarable/platform/mkldnn/lrn.cpp
+96
View File
@@ -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
libnd4j/include/ops/declarable/platform/mkldnn/lstmLayer.cpp
+547
View File
@@ -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
libnd4j/include/ops/declarable/platform/mkldnn/matmul.cpp
+328
View File
@@ -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
libnd4j/include/ops/declarable/platform/mkldnn/maxpooling2d.cpp
+148
View File
@@ -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
libnd4j/include/ops/declarable/platform/mkldnn/maxpooling3d.cpp
+156
View File
@@ -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
libnd4j/include/ops/declarable/platform/mkldnn/mkldnnUtils.cpp
+440
View File
@@ -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
libnd4j/include/ops/declarable/platform/mkldnn/mkldnnUtils.h
+171
View File
@@ -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
libnd4j/include/ops/declarable/platform/mkldnn/softmax.cpp
+268
View File
@@ -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
libnd4j/include/ops/declarable/platform/mkldnn/tanh.cpp
+218
View File
@@ -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
libnd4j/include/ops/declarable/platform/mkldnn/xw_plus_b.cpp
+415
View File
@@ -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