cavis/libnd4j/include/ops/declarable/helpers/cpu/addBias.cpp

159 lines
7.1 KiB
C++

/*******************************************************************************
* Copyright (c) 2015-2018 Skymind, Inc.
*
* 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 Yurii Shyrma, created on 26.02.2018
//
#include<ops/declarable/helpers/addBias.h>
#include <execution/Threads.h>
namespace nd4j {
namespace ops {
namespace helpers {
//////////////////////////////////////////////////////////////////////////
template <typename X, typename Y>
static void addBias_(const NDArray& input, const NDArray& bias, NDArray &output, const bool isNCHW) {
// bias [oC]
// if(input_rank == 4)
// input and output have same shapes: [bS, oH, oW, oC] (NHWC) or [bS, oC, oH, oW] (NCHW)
// if(input_rank == 5)
// input and output have same shapes: [bS, oD, oH, oW, oC] (NHWC) or [bS, oD, oC, oH, oW] (NCHW)
// else
// apply applyBroadCast
const X* x = input.bufferAsT<X>();
const Y* y = bias.bufferAsT<Y>();
X* z = output.bufferAsT<X>();
const bool inOutAreSame = x == z;
const uint bS = output.sizeAt(0); // batch size
const Nd4jLong yStrideC = bias.stridesOf()[0];
const Nd4jLong zStrideB = output.stridesOf()[0];
if(output.rankOf() == 4) {
const uint C = isNCHW ? output.sizeAt(1) : output.sizeAt(3); // channels
const uint oH = isNCHW ? output.sizeAt(2) : output.sizeAt(1); // height
const uint oW = isNCHW ? output.sizeAt(3) : output.sizeAt(2); // width
const Nd4jLong zStrideC = isNCHW ? output.stridesOf()[1] : output.stridesOf()[3];
const Nd4jLong zStrideH = isNCHW ? output.stridesOf()[2] : output.stridesOf()[1];
const Nd4jLong zStrideW = isNCHW ? output.stridesOf()[3] : output.stridesOf()[2];
if(inOutAreSame) {
auto func = PRAGMA_THREADS_FOR_3D {
for (uint b = start_x; b < stop_x; b += inc_x)
for (uint c = start_y; c < stop_y; c += inc_y)
for (uint h = start_z; h < stop_z; h += inc_z)
for (uint w = 0; w < oW; ++w)
z[b * zStrideB + c * zStrideC + h * zStrideH + w * zStrideW] += static_cast<X>(y[c * yStrideC]);
};
samediff::Threads::parallel_for(func, 0, bS, 1, 0, C, 1, 0, oH, 1);
}
else {
const Nd4jLong xStrideB = input.stridesOf()[0];
const Nd4jLong xStrideC = isNCHW ? input.stridesOf()[1] : input.stridesOf()[3];
const Nd4jLong xStrideH = isNCHW ? input.stridesOf()[2] : input.stridesOf()[1];
const Nd4jLong xStrideW = isNCHW ? input.stridesOf()[3] : input.stridesOf()[2];
auto func = PRAGMA_THREADS_FOR_3D {
for (uint b = start_x; b < stop_x; b += inc_x)
for (uint c = start_y; c < stop_y; c += inc_y)
for (uint h = start_z; h < stop_z; h += inc_z)
for (uint w = 0; w < oW; ++w)
z[b * zStrideB + c * zStrideC + h * zStrideH + w * zStrideW] = x[b * xStrideB + c * xStrideC + h * xStrideH + w * xStrideW] + static_cast<X>(y[c * yStrideC]);
};
samediff::Threads::parallel_for(func, 0, bS, 1, 0, C, 1, 0, oH, 1);
}
}
else if(output.rankOf() == 5) {
const uint C = isNCHW ? output.sizeAt(1) : output.sizeAt(4); // channels
const uint oD = isNCHW ? output.sizeAt(2) : output.sizeAt(1); // depth
const uint oH = isNCHW ? output.sizeAt(3) : output.sizeAt(2); // height
const uint oW = isNCHW ? output.sizeAt(4) : output.sizeAt(3); // width
const Nd4jLong zStrideC = isNCHW ? output.stridesOf()[1] : output.stridesOf()[4];
const Nd4jLong zStrideD = isNCHW ? output.stridesOf()[2] : output.stridesOf()[1];
const Nd4jLong zStrideH = isNCHW ? output.stridesOf()[3] : output.stridesOf()[2];
const Nd4jLong zStrideW = isNCHW ? output.stridesOf()[4] : output.stridesOf()[3];
if(inOutAreSame) {
auto func = PRAGMA_THREADS_FOR_3D {
for (uint b = start_x; b < stop_x; b += inc_x)
for (uint c = start_y; c < stop_y; c += inc_y)
for (uint d = start_z; d < stop_z; d += inc_z)
for (uint h = 0; h < oH; ++h)
for (uint w = 0; w < oW; ++w)
z[b * zStrideB + c * zStrideC + d * zStrideD + h * zStrideH + w * zStrideW] += static_cast<X>(y[c * yStrideC]);
};
samediff::Threads::parallel_for(func, 0, bS, 1, 0, C, 1, 0, oD, 1);
}
else {
const Nd4jLong xStrideB = input.stridesOf()[0];
const Nd4jLong xStrideC = isNCHW ? input.stridesOf()[1] : input.stridesOf()[4];
const Nd4jLong xStrideD = isNCHW ? input.stridesOf()[2] : input.stridesOf()[1];
const Nd4jLong xStrideH = isNCHW ? input.stridesOf()[3] : input.stridesOf()[2];
const Nd4jLong xStrideW = isNCHW ? input.stridesOf()[4] : input.stridesOf()[3];
auto func = PRAGMA_THREADS_FOR_3D {
for (uint b = start_x; b < stop_x; b += inc_x)
for (uint c = start_y; c < stop_y; c += inc_y)
for (uint d = start_z; d < stop_z; d += inc_z)
for (uint h = 0; h < oH; ++h)
for (uint w = 0; w < oW; ++w)
z[b * zStrideB + c * zStrideC + d * zStrideD + h * zStrideH + w * zStrideW] = x[b * xStrideB + c * xStrideC + d * xStrideD + h * xStrideH + w * xStrideW] + static_cast<X>(y[c * yStrideC]);
};
samediff::Threads::parallel_for(func, 0, bS, 1, 0, C, 1, 0, oD, 1);
}
}
else {
const int channelDim = isNCHW ? 1 : input.rankOf() - 1; // second or last
const_cast<NDArray&>(input).applyBroadcast(nd4j::broadcast::Add, {channelDim}, &bias, &output);
}
}
//////////////////////////////////////////////////////////////////////////
void addBias(nd4j::graph::Context& block, const NDArray& input, const NDArray& bias, NDArray& output, const bool isNCHW) {
// bias.rankOf() == 1 ? bias : bias.reshape(bias.ordering(), {bias.lengthOf()})
BUILD_DOUBLE_SELECTOR(input.dataType(), bias.dataType(), addBias_, (input, bias, output, isNCHW), FLOAT_TYPES, FLOAT_TYPES);
}
BUILD_DOUBLE_TEMPLATE(template void addBias_, (const NDArray& input, const NDArray& bias, NDArray& output, const bool isNCHW), FLOAT_TYPES, FLOAT_TYPES);
}
}
}