wehub-resource-sync
627 lines
26 KiB
Plaintext
627 lines
26 KiB
Plaintext
/* ******************************************************************************
|
|
*
|
|
*
|
|
* 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
|
|
******************************************************************************/
|
|
|
|
//
|
|
// implementation of operations for Simple Recurrent Unit: arXiv:1709.02755v2 [cs.CL] 12 Sep 2017
|
|
//
|
|
// @author Yurii Shyrma, created on 05.12.2017
|
|
//
|
|
#include <array/NDArrayFactory.h>
|
|
#include <helpers/MmulHelper.h>
|
|
#include <helpers/PointersManager.h>
|
|
#include <ops/declarable/helpers/sru.h>
|
|
|
|
#include "execution/cuda/LaunchDims.h"
|
|
|
|
|
|
namespace sd {
|
|
namespace ops {
|
|
namespace helpers {
|
|
|
|
//////////////////////////////////////////////////////////////////////////
|
|
static SD_INLINE NDArray activation(NDArray& arr) {
|
|
auto result = NDArray(&arr, false, arr.getContext());
|
|
arr.applyTransform(transform::Tanh,&result);
|
|
return result;
|
|
}
|
|
|
|
//////////////////////////////////////////////////////////////////////////
|
|
static SD_INLINE NDArray sigmoid(NDArray& arr) {
|
|
return (const_cast<NDArray&>(arr)).transform(transform::Sigmoid);
|
|
}
|
|
|
|
//////////////////////////////////////////////////////////////////////////
|
|
void sruCell(LaunchContext* context, NDArray* x, NDArray* c0, NDArray* w, NDArray* b,
|
|
NDArray* h, NDArray* c) {
|
|
// x input [bS x inSize], bS - batch size, inSize - number of features
|
|
// c0 previous cell state c [bS x inSize], that is at previous time step t-1
|
|
// w weights [inSize x 3*inSize]
|
|
// b biases [2*inSize]
|
|
|
|
// h current cell output [bS x inSize], that is at current time step t
|
|
// c current cell state [bS x inSize], that is at current time step t
|
|
|
|
const int inSize = x->sizeAt(1); // inSize - number of features
|
|
|
|
auto z = mmul(*x, *w); // [bS x 3*inSize]
|
|
|
|
// forget gate = sigmoid(x*Wf + bf)
|
|
NDArray fIn = z({0, 0, inSize, 2 * inSize}) + (*b)({0, inSize});
|
|
auto f = sigmoid(fIn);
|
|
|
|
NDArray rIn = z({0, 0, 2 * inSize, 3 * inSize}) + (*b)({inSize, 2 * inSize});
|
|
// reset gate = sigmoid(x*Wr + br)
|
|
auto r = sigmoid(rIn);
|
|
|
|
// ◦ means element-wise product or so called Hadamard product
|
|
// current sell state = f◦c0 + (1 - f)◦(x*Wc)
|
|
NDArray cAssign = f * (*c0) + (1.f - f) * z({0, 0, 0, inSize});
|
|
c->assign(&cAssign);
|
|
// *c = f*(*c0 - z({},{0, inSize})) + z({{},{0, inSize}});
|
|
|
|
// current cell output = r◦activation(c) + (1 - r)◦x
|
|
NDArray resultTwo = r * activation(*c) + (1.f - r) * (*x);
|
|
h->assign(&resultTwo);
|
|
// *h = r * (activation<T>(c) - *x) + *x;
|
|
}
|
|
|
|
//////////////////////////////////////////////////////////////////////////
|
|
void sruTimeLoop(LaunchContext* context, NDArray* x, NDArray* c0, NDArray* w, NDArray* b,
|
|
NDArray* h, NDArray* c) {
|
|
// x input [bS x inSize x time]
|
|
// c0 initial cell state (at time step = 0) [bS x inSize],
|
|
// w weights, [3*inSize x inSize]
|
|
// b biases, [2*inSize]
|
|
|
|
// h cell outputs [bS x inSize x time]
|
|
// c cell states [bS x inSize x time]
|
|
|
|
auto wT = w->transpose(); // [3*inSize x inSize] -> [inSize x 3*inSize]
|
|
|
|
const int time = x->sizeAt(2);
|
|
|
|
NDArray ct_1(*c0);
|
|
|
|
// loop through time steps
|
|
for (int t = 0; t < time; ++t) {
|
|
auto xt = (*x)({0, 0, 0, 0, t, t + 1});
|
|
auto ht = (*h)({0, 0, 0, 0, t, t + 1});
|
|
auto ct = (*c)({0, 0, 0, 0, t, t + 1});
|
|
|
|
sruCell(context, &xt, &ct_1, &wT, b, &ht, &ct);
|
|
ct_1.assign(&ct);
|
|
}
|
|
|
|
delete wT;
|
|
}
|
|
|
|
//////////////////////////////////////////////////////////////////////////
|
|
template <typename T>
|
|
__global__ static void sruBICuda(const void* vx, const LongType* xShapeInfo, const void* vwi,
|
|
const LongType* wiShapeInfo, const void* vb, const LongType* bShapeInfo,
|
|
const void* vc0, const LongType* c0ShapeInfo, const void* vmask,
|
|
const LongType* maskShapeInfo, void* vht, const LongType* htShapeInfo,
|
|
void* vct, const LongType* ctShapeInfo) {
|
|
// Inputs:
|
|
// x [time, bS, 2*K]
|
|
// wi [time, bS, 6*K], wi = mmul(x, weights);
|
|
// b [4*K]
|
|
// c0 [bS, 2*K]
|
|
// mask [bS, 2*K], optional
|
|
|
|
// Outputs:
|
|
// ht [time, bS, 2*K]
|
|
// ct [time, bS, 2*K]
|
|
|
|
// Reinterpret inputs and outputs
|
|
const T* x = reinterpret_cast<const T*>(vx);
|
|
const T* wi = reinterpret_cast<const T*>(vwi);
|
|
const T* b = reinterpret_cast<const T*>(vb);
|
|
const T* c0 = reinterpret_cast<const T*>(vc0);
|
|
const T* mask = reinterpret_cast<const T*>(vmask);
|
|
T* ht = reinterpret_cast<T*>(vht);
|
|
T* ct = reinterpret_cast<T*>(vct);
|
|
|
|
const int rank = 3; // Assuming 3D tensors
|
|
|
|
// Shared memory for caching shape information and other variables
|
|
extern __shared__ unsigned char shmem[];
|
|
// Pointers to shared memory segments
|
|
LongType* sharedMem = reinterpret_cast<LongType*>(shmem);
|
|
|
|
// Shared variables
|
|
__shared__ LongType shared_time;
|
|
__shared__ LongType shared_bS;
|
|
__shared__ LongType shared_K;
|
|
__shared__ LongType shared_len;
|
|
__shared__ LongType shared_totalThreads;
|
|
|
|
// Cached shape and stride pointers
|
|
__shared__ const LongType* shared_xShape;
|
|
__shared__ const LongType* shared_wiShape;
|
|
__shared__ const LongType* shared_bShape;
|
|
__shared__ const LongType* shared_c0Shape;
|
|
__shared__ const LongType* shared_maskShape;
|
|
__shared__ const LongType* shared_htShape;
|
|
__shared__ const LongType* shared_ctShape;
|
|
|
|
if (threadIdx.x == 0) {
|
|
// Cache shape pointers
|
|
shared_xShape = shape::shapeOf(xShapeInfo);
|
|
shared_wiShape = shape::shapeOf(wiShapeInfo);
|
|
shared_bShape = shape::shapeOf(bShapeInfo);
|
|
shared_c0Shape = shape::shapeOf(c0ShapeInfo);
|
|
shared_maskShape = shape::shapeOf(maskShapeInfo);
|
|
shared_htShape = shape::shapeOf(htShapeInfo);
|
|
shared_ctShape = shape::shapeOf(ctShapeInfo);
|
|
|
|
// Cache time, bS, and K
|
|
shared_time = shared_xShape[0]; // time
|
|
shared_bS = shared_xShape[1]; // batch size (bS)
|
|
shared_K = shared_xShape[2] / 2; // Assuming xShapeInfo[2] = 2*K
|
|
|
|
// Calculate len = 2*K * bS
|
|
shared_len = 2 * shared_K * shared_bS;
|
|
|
|
// Calculate total number of threads across all blocks
|
|
shared_totalThreads = gridDim.x * blockDim.x;
|
|
}
|
|
|
|
// Ensure all threads have access to the cached values
|
|
__syncthreads();
|
|
|
|
// Calculate the global thread ID
|
|
const LongType tid = blockIdx.x * blockDim.x + threadIdx.x;
|
|
|
|
// Allocate space in shared memory for coordinates
|
|
LongType* coords = sharedMem + threadIdx.x * (rank - 1); // Only last two dimensions {bS, 2*K}
|
|
|
|
if (tid >= shared_len) return;
|
|
|
|
// Convert linear index to multi-dimensional coordinates {bS, 2*K}
|
|
INDEX2COORDS(tid, rank - 1, shared_xShape, coords); // coords[0] = bS, coords[1] = 2*K
|
|
|
|
// Calculate necessary offsets
|
|
LongType maskOffset = 0, c0Offset = 0, bFOffset = 0, bROffset = 0;
|
|
|
|
if (vmask != nullptr) {
|
|
COORDS2INDEX(rank - 1, shape::stride(maskShapeInfo), coords, maskOffset);
|
|
}
|
|
COORDS2INDEX(rank - 1, shape::stride(c0ShapeInfo), coords, c0Offset);
|
|
COORDS2INDEX(rank - 1, shape::stride(bShapeInfo), coords + 1, bFOffset);
|
|
bROffset = bFOffset + 2 * shared_K * shared_bShape[2]; // 2*K*b_stride
|
|
|
|
// Fetch values
|
|
const T maskVal = (vmask != nullptr) ? mask[maskOffset] : static_cast<T>(1);
|
|
const T bF = b[bFOffset];
|
|
const T bR = b[bROffset];
|
|
T c0Val = c0[c0Offset];
|
|
|
|
// Determine flip condition
|
|
const bool flip = coords[1] >= shared_K;
|
|
|
|
// Initialize coordinates for time iteration
|
|
if (flip)
|
|
coords[0] = shared_time - 1;
|
|
else
|
|
coords[0] = 0;
|
|
|
|
// Calculate offsets for x, ht, ct
|
|
LongType xOffset = 0, htOffset = 0, ctOffset = 0;
|
|
COORDS2INDEX(rank, shape::stride(xShapeInfo), coords, xOffset);
|
|
COORDS2INDEX(rank, shape::stride(htShapeInfo), coords, htOffset);
|
|
COORDS2INDEX(rank, shape::stride(ctShapeInfo), coords, ctOffset);
|
|
|
|
// Adjust coords for wi and gradWi
|
|
coords[1] *= 3; // 6*K corresponds to 3 * 2*K
|
|
|
|
// Calculate wi offsets
|
|
LongType wiOffset0 = 0, wiOffset1 = 0, wiOffset2 = 0;
|
|
COORDS2INDEX(rank, shape::stride(wiShapeInfo), coords, wiOffset0);
|
|
wiOffset1 = wiOffset0 + shared_wiShape[rank]; // Add stride for wi1
|
|
wiOffset2 = wiOffset1 + shared_wiShape[rank]; // Add stride for wi2
|
|
|
|
// Iterate over the time steps
|
|
for (LongType t = 0; t < shared_time; ++t) {
|
|
// Evaluate sigmoids
|
|
T ft = static_cast<T>(1) / (static_cast<T>(1) + math::sd_exp<T, T>(- (wi[wiOffset1] + bF)));
|
|
T rt = static_cast<T>(1) / (static_cast<T>(1) + math::sd_exp<T, T>(- (wi[wiOffset2] + bR)));
|
|
|
|
// Update c0Val and ct
|
|
c0Val = (c0Val - wi[wiOffset0]) * ft + wi[wiOffset0];
|
|
ct[ctOffset] = c0Val;
|
|
|
|
// Compute tanh activation
|
|
T val = math::sd_tanh<T, T>(c0Val);
|
|
|
|
// Fetch x value
|
|
T xVal = x[xOffset];
|
|
|
|
// Compute ht
|
|
ht[htOffset] = (val * maskVal - xVal) * rt + xVal;
|
|
|
|
// Update offsets based on flip condition
|
|
if (flip) {
|
|
xOffset -= shape::stride(xShapeInfo)[0]; // time step stride
|
|
htOffset -= shape::stride(htShapeInfo)[0];
|
|
ctOffset -= shape::stride(ctShapeInfo)[0];
|
|
wiOffset0 -= shape::stride(wiShapeInfo)[0];
|
|
wiOffset1 -= shape::stride(wiShapeInfo)[0];
|
|
wiOffset2 -= shape::stride(wiShapeInfo)[0];
|
|
} else {
|
|
xOffset += shape::stride(xShapeInfo)[0]; // time step stride
|
|
htOffset += shape::stride(htShapeInfo)[0];
|
|
ctOffset += shape::stride(ctShapeInfo)[0];
|
|
wiOffset0 += shape::stride(wiShapeInfo)[0];
|
|
wiOffset1 += shape::stride(wiShapeInfo)[0];
|
|
wiOffset2 += shape::stride(wiShapeInfo)[0];
|
|
}
|
|
}
|
|
}
|
|
|
|
//////////////////////////////////////////////////////////////////////////
|
|
template <typename T>
|
|
static void sruBICudaLauncher(const int blocksPerGrid, const int threadsPerBlock, const int sharedMem,
|
|
const cudaStream_t* stream, const void* vx, const LongType* xShapeInfo,
|
|
const void* vwi,
|
|
const LongType* wiShapeInfo, const void* vb, const LongType* bShapeInfo, const void* vc0,
|
|
const LongType* c0ShapeInfo,
|
|
const void* vmask, const LongType* maskShapeInfo, void* vht,
|
|
const LongType* htShapeInfo, void* vct, const LongType* ctShapeInfo) {
|
|
sruBICuda<T><<<blocksPerGrid, threadsPerBlock, sharedMem, *stream>>>(vx, xShapeInfo, vwi, wiShapeInfo, vb, bShapeInfo,
|
|
vc0, c0ShapeInfo, vmask, maskShapeInfo, vht,
|
|
htShapeInfo, vct, ctShapeInfo);
|
|
sd::DebugHelper::checkErrorCode(const_cast<cudaStream_t *>(stream), "sruBICuda failed");
|
|
|
|
}
|
|
|
|
//////////////////////////////////////////////////////////////////////////
|
|
void sruBI(LaunchContext* context, NDArray* x, NDArray* w, NDArray* b, NDArray* c0,
|
|
NDArray* mask, NDArray* ht, NDArray* ct) {
|
|
// x = x * mask
|
|
std::vector<LongType> dims = {1,2};
|
|
if (mask) x->applyBroadcast(broadcast::Multiply, &dims, mask, x); // apply mask
|
|
|
|
// U = x * w
|
|
NDArray wi = mmul(*x, *w); // U [time x bS x 6*K]
|
|
|
|
PointersManager manager(context, "sru_bi");
|
|
|
|
dim3 sruBiDims2 = sruBiDims(x->sizeAt(1) * x->sizeAt(2),x->rankOf());
|
|
NDArray::prepareSpecialUse({ht, ct}, {x, &wi, b, c0, mask});
|
|
BUILD_SINGLE_SELECTOR(
|
|
x->dataType(), sruBICudaLauncher,
|
|
(sruBiDims2.y,sruBiDims2.x, sruBiDims2.z, context->getCudaStream(), x->specialBuffer(), x->specialShapeInfo(),
|
|
wi.specialBuffer(), wi.specialShapeInfo(), b->specialBuffer(), b->specialShapeInfo(), c0->specialBuffer(),
|
|
c0->specialShapeInfo(), mask ? mask->specialBuffer() : nullptr, mask ? mask->specialShapeInfo() : nullptr,
|
|
ht->specialBuffer(), ht->specialShapeInfo(), ct->specialBuffer(), ct->specialShapeInfo()),
|
|
SD_FLOAT_TYPES);
|
|
NDArray::registerSpecialUse({ht, ct}, {x, &wi, b, c0, mask});
|
|
|
|
manager.synchronize();
|
|
}
|
|
|
|
//////////////////////////////////////////////////////////////////////////
|
|
template <typename T>
|
|
__global__ static void sruBIBPCuda(const void* vx, const LongType* xShapeInfo, const void* vwi,
|
|
const LongType* wiShapeInfo, const void* vb, const LongType* bShapeInfo,
|
|
const void* vc0, const LongType* c0ShapeInfo, const void* vmask,
|
|
const LongType* maskShapeInfo, const void* vct, const LongType* ctShapeInfo,
|
|
const void* vgradHt, const LongType* gradHtShapeInfo, const void* vgradCt,
|
|
const LongType* gradCtShapeInfo, void* vgradI, const LongType* gradIShapeInfo,
|
|
void* vgradWi, const LongType* gradWiShapeInfo, void* vgradB,
|
|
const LongType* gradBShapeInfo, void* vgradC0, const LongType* gradC0ShapeInfo) {
|
|
// Inputs:
|
|
// x [time, bS, 2*K]
|
|
// wi [time, bS, 6*K], wi = mmul(x, weights);
|
|
// b [4*K]
|
|
// c0 [bS, 2*K]
|
|
// mask [bS, 2*K], optional
|
|
// ct [time, bS, 2*K]
|
|
// gradHt [time, bS, 2*K]
|
|
// gradCt [bS, 2*K]
|
|
|
|
// Outputs:
|
|
// gradI [time, bS, 2*K]
|
|
// gradWi [time, 2*K, 6*K]
|
|
// gradB [bS, 4*K]
|
|
// gradC0 [bS, 2*K]
|
|
|
|
// Reinterpret inputs and outputs
|
|
const T* x = reinterpret_cast<const T*>(vx);
|
|
const T* wi = reinterpret_cast<const T*>(vwi);
|
|
const T* b = reinterpret_cast<const T*>(vb);
|
|
const T* c0 = reinterpret_cast<const T*>(vc0);
|
|
const T* mask = reinterpret_cast<const T*>(vmask);
|
|
const T* ct = reinterpret_cast<const T*>(vct);
|
|
const T* gradHt = reinterpret_cast<const T*>(vgradHt);
|
|
const T* gradCt = reinterpret_cast<const T*>(vgradCt);
|
|
|
|
T* gradI = reinterpret_cast<T*>(vgradI);
|
|
T* gradWi = reinterpret_cast<T*>(vgradWi);
|
|
T* gradB = reinterpret_cast<T*>(vgradB);
|
|
T* gradC0 = reinterpret_cast<T*>(vgradC0);
|
|
|
|
const int rank = 3; // Assuming 3D tensors
|
|
|
|
// Shared memory for caching shape information
|
|
extern __shared__ unsigned char shmem[];
|
|
LongType* sharedMem = reinterpret_cast<LongType*>(shmem);
|
|
|
|
__shared__ LongType shared_time;
|
|
__shared__ LongType shared_K;
|
|
__shared__ LongType shared_len;
|
|
__shared__ LongType shared_totalThreads;
|
|
|
|
// Cached shape pointers
|
|
__shared__ const LongType* shared_xShape;
|
|
__shared__ const LongType* shared_wiShape;
|
|
__shared__ const LongType* shared_bShape;
|
|
__shared__ const LongType* shared_c0Shape;
|
|
__shared__ const LongType* shared_maskShape;
|
|
__shared__ const LongType* shared_ctShape;
|
|
__shared__ const LongType* shared_gradHtShape;
|
|
__shared__ const LongType* shared_gradCtShape;
|
|
__shared__ const LongType* shared_gradIShape;
|
|
__shared__ const LongType* shared_gradWiShape;
|
|
__shared__ const LongType* shared_gradBShape;
|
|
__shared__ const LongType* shared_gradC0Shape;
|
|
|
|
if (threadIdx.x == 0) {
|
|
// Cache ranks, shapes, and strides
|
|
shared_xShape = shape::shapeOf(xShapeInfo);
|
|
shared_wiShape = shape::shapeOf(wiShapeInfo);
|
|
shared_bShape = shape::shapeOf(bShapeInfo);
|
|
shared_c0Shape = shape::shapeOf(c0ShapeInfo);
|
|
shared_maskShape = shape::shapeOf(maskShapeInfo);
|
|
shared_ctShape = shape::shapeOf(ctShapeInfo);
|
|
shared_gradHtShape = shape::shapeOf(gradHtShapeInfo);
|
|
shared_gradCtShape = shape::shapeOf(gradCtShapeInfo);
|
|
shared_gradIShape = shape::shapeOf(gradIShapeInfo);
|
|
shared_gradWiShape = shape::shapeOf(gradWiShapeInfo);
|
|
shared_gradBShape = shape::shapeOf(gradBShapeInfo);
|
|
shared_gradC0Shape = shape::shapeOf(gradC0ShapeInfo);
|
|
|
|
// Cache time and K
|
|
shared_time = shared_xShape[0];
|
|
shared_K = shared_xShape[2] / 2; // Assuming xShapeInfo[2] = 2*K
|
|
|
|
// Calculate len = 2*K * bS
|
|
LongType bS = shared_xShape[1];
|
|
shared_len = 2 * shared_K * bS;
|
|
|
|
// Total threads across all blocks
|
|
shared_totalThreads = gridDim.x * blockDim.x;
|
|
}
|
|
|
|
// Ensure all threads have access to the cached values
|
|
__syncthreads();
|
|
|
|
const LongType tid = blockIdx.x * blockDim.x + threadIdx.x;
|
|
|
|
// Allocate space in shared memory for coordinates
|
|
LongType* coords = sharedMem + threadIdx.x * rank;
|
|
|
|
if (tid >= shared_len) return;
|
|
|
|
// Convert linear index to coordinates {bS, 2*K}
|
|
INDEX2COORDS(tid, rank - 1, shared_xShape, coords + 1); // Skipping the time dimension
|
|
|
|
// Calculate necessary offsets
|
|
LongType maskOffset = 0, c0Offset = 0, gradCtOffset = 0, gradC0Offset = 0;
|
|
LongType bFOffset = 0, bROffset = 0, gradBFOffset = 0, gradBROffset = 0;
|
|
|
|
if (vmask != nullptr) {
|
|
COORDS2INDEX(rank - 1, shape::stride(maskShapeInfo), coords + 1, maskOffset);
|
|
}
|
|
COORDS2INDEX(rank - 1, shape::stride(c0ShapeInfo), coords + 1, c0Offset);
|
|
COORDS2INDEX(rank - 1, shape::stride(gradCtShapeInfo), coords + 1, gradCtOffset);
|
|
COORDS2INDEX(rank - 1, shape::stride(gradC0ShapeInfo), coords + 1, gradC0Offset);
|
|
COORDS2INDEX(rank - 1, shape::stride(bShapeInfo), coords + 2, bFOffset);
|
|
bROffset = bFOffset + 2 * shared_K * shared_bShape[2]; // 2*K*b_stride
|
|
gradBFOffset = coords[1] * shared_gradBShape[3] / 2 + coords[2] * shared_gradBShape[4];
|
|
gradBROffset = gradBFOffset + shared_gradBShape[3];
|
|
|
|
const bool flip = coords[2] >= shared_K;
|
|
|
|
if (flip)
|
|
coords[0] = 0;
|
|
else
|
|
coords[0] = shared_time - 1;
|
|
|
|
// Calculate offsets for x, ct, gradI, gradHt
|
|
LongType xOffset = 0, ctOffset = 0, gradIOffset = 0, gradHtOffset = 0;
|
|
COORDS2INDEX(rank, shape::stride(xShapeInfo), coords, xOffset);
|
|
COORDS2INDEX(rank, shape::stride(ctShapeInfo), coords, ctOffset);
|
|
COORDS2INDEX(rank, shape::stride(gradIShapeInfo), coords, gradIOffset);
|
|
COORDS2INDEX(rank, shape::stride(gradHtShapeInfo), coords, gradHtOffset);
|
|
|
|
// Adjust coords for wi and gradWi
|
|
coords[2] *= 3;
|
|
LongType gradWiOffset0 = 0, gradWiOffset1 = 0, gradWiOffset2 = 0;
|
|
LongType wiOffset0 = 0, wiOffset1 = 0, wiOffset2 = 0;
|
|
|
|
COORDS2INDEX(rank, shape::stride(gradWiShapeInfo), coords, gradWiOffset0);
|
|
gradWiOffset1 = gradWiOffset0 + shared_gradWiShape[rank + 3]; // add last stride
|
|
gradWiOffset2 = gradWiOffset1 + shared_gradWiShape[rank + 3]; // add last stride
|
|
|
|
COORDS2INDEX(rank, shape::stride(wiShapeInfo), coords, wiOffset0);
|
|
wiOffset1 = wiOffset0 + shared_wiShape[rank + 3]; // add last stride
|
|
wiOffset2 = wiOffset1 + shared_wiShape[rank + 3]; // add last stride
|
|
|
|
// Fetch values
|
|
const T xVal = x[xOffset];
|
|
const T maskVal = (vmask != nullptr) ? mask[maskOffset] : static_cast<T>(1);
|
|
const T c0Val = c0[c0Offset];
|
|
const T bF = b[bFOffset];
|
|
const T bR = b[bROffset];
|
|
T gradCtVal = gradCt[gradCtOffset];
|
|
T gbF = static_cast<T>(0);
|
|
T gbR = static_cast<T>(0);
|
|
|
|
// Iterate over the time steps
|
|
for (LongType t = 0; t < shared_time; ++t) {
|
|
// Evaluate sigmoids
|
|
T ft = static_cast<T>(1) / (static_cast<T>(1) + math::sd_exp<T, T>(- (wi[wiOffset1] + bF)));
|
|
T rt = static_cast<T>(1) / (static_cast<T>(1) + math::sd_exp<T, T>(- (wi[wiOffset2] + bR)));
|
|
|
|
T val = math::sd_tanh<T, T>(ct[ctOffset]);
|
|
|
|
T prevVal;
|
|
if (t < shared_time - 1)
|
|
prevVal = ct[ctOffset += (flip ? shared_ctShape[rank + 1] : -shared_ctShape[rank + 1])];
|
|
else
|
|
prevVal = c0Val;
|
|
|
|
// Gradient with respect to input
|
|
gradI[gradIOffset] = gradHt[gradHtOffset] - gradHt[gradHtOffset] * rt;
|
|
|
|
// Gradient with respect to rt, wiR, and bR
|
|
T grt = gradHt[gradHtOffset] * (val * maskVal - x[xOffset]) * (rt - rt * rt);
|
|
gradWi[gradWiOffset2] = grt;
|
|
gbR += grt;
|
|
|
|
// Gradient with respect to state
|
|
T gradC0Val = gradHt[gradHtOffset] * maskVal * (rt - rt * val * val) + gradCtVal;
|
|
|
|
// Gradient with respect to wi0
|
|
gradWi[gradWiOffset0] = gradC0Val - gradC0Val * ft;
|
|
|
|
// Gradient with respect to ft, wi1, and bF
|
|
T gft = gradC0Val * (prevVal - wi[wiOffset0]) * (ft - ft * ft);
|
|
gradWi[gradWiOffset1] = gft;
|
|
gbF += gft;
|
|
|
|
// Gradient with respect to c_previous
|
|
gradCtVal = gradC0Val * ft;
|
|
|
|
// Update offsets based on flip
|
|
if (flip) {
|
|
xOffset += shared_xShape[rank + 1]; // first stride, corresponds to time step
|
|
gradHtOffset += shared_gradHtShape[rank + 1];
|
|
gradIOffset += shared_gradIShape[rank + 1];
|
|
wiOffset0 += shared_wiShape[rank + 1];
|
|
wiOffset1 += shared_wiShape[rank + 1];
|
|
wiOffset2 += shared_wiShape[rank + 1];
|
|
gradWiOffset0 += shared_gradWiShape[rank + 1];
|
|
gradWiOffset1 += shared_gradWiShape[rank + 1];
|
|
gradWiOffset2 += shared_gradWiShape[rank + 1];
|
|
}
|
|
else {
|
|
xOffset -= shared_xShape[rank + 1]; // first stride, corresponds to time step
|
|
gradHtOffset -= shared_gradHtShape[rank + 1];
|
|
gradIOffset -= shared_gradIShape[rank + 1];
|
|
wiOffset0 -= shared_wiShape[rank + 1];
|
|
wiOffset1 -= shared_wiShape[rank + 1];
|
|
wiOffset2 -= shared_wiShape[rank + 1];
|
|
gradWiOffset0 -= shared_gradWiShape[rank + 1];
|
|
gradWiOffset1 -= shared_gradWiShape[rank + 1];
|
|
gradWiOffset2 -= shared_gradWiShape[rank + 1];
|
|
}
|
|
}
|
|
|
|
// Write accumulated gradients to output
|
|
gradB[gradBFOffset] = gbF;
|
|
gradB[gradBROffset] = gbR;
|
|
gradC0[gradC0Offset] = gradCtVal;
|
|
}
|
|
|
|
|
|
//////////////////////////////////////////////////////////////////////////
|
|
template <typename T>
|
|
static void sruBIBPCudaLauncher(
|
|
const int blocksPerGrid, const int threadsPerBlock, const int sharedMem, const cudaStream_t* stream, const void* vx, const LongType* xShapeInfo, const void* vwi,
|
|
const LongType* wiShapeInfo, const void* vb, const LongType* bShapeInfo, const void* vc0, const LongType* c0ShapeInfo, const void* vmask,
|
|
const LongType* maskShapeInfo, const void* vct, const LongType* ctShapeInfo, const void* vgradHt, const LongType* gradHtShapeInfo, const void* vgradCt,
|
|
const LongType* gradCtShapeInfo, void* vgradI, const LongType* gradIShapeInfo, void* vgradWi, const LongType* gradWiShapeInfo, void* vgradB,
|
|
const LongType* gradBShapeInfo, void* vgradC0, const LongType* gradC0ShapeInfo) {
|
|
sruBIBPCuda<T><<<blocksPerGrid, threadsPerBlock, sharedMem, *stream>>>(
|
|
vx, xShapeInfo, vwi, wiShapeInfo, vb, bShapeInfo, vc0, c0ShapeInfo, vmask, maskShapeInfo, vct, ctShapeInfo,
|
|
vgradHt, gradHtShapeInfo, vgradCt, gradCtShapeInfo, vgradI, gradIShapeInfo, vgradWi, gradWiShapeInfo, vgradB,
|
|
gradBShapeInfo, vgradC0, gradC0ShapeInfo);
|
|
sd::DebugHelper::checkErrorCode(const_cast<cudaStream_t *>(stream), "sruBIBPCuda failed");
|
|
|
|
}
|
|
BUILD_SINGLE_TEMPLATE( void sruBIBPCudaLauncher,
|
|
(const int blocksPerGrid, const int threadsPerBlock, const int sharedMem,
|
|
const cudaStream_t* stream, const void* vx, const sd::LongType* xShapeInfo, const void* vwi,
|
|
const sd::LongType* wiShapeInfo, const void* vb, const sd::LongType* bShapeInfo, const void* vc0,
|
|
const sd::LongType* c0ShapeInfo, const void* vmask, const sd::LongType* maskShapeInfo,
|
|
const void* vct, const sd::LongType* ctShapeInfo, const void* vgradHt,
|
|
const sd::LongType* gradHtShapeInfo, const void* vgradCt, const sd::LongType* gradCtShapeInfo,
|
|
void* vgradI, const sd::LongType* gradIShapeInfo, void* vgradWi,
|
|
const sd::LongType* gradWiShapeInfo, void* vgradB, const sd::LongType* gradBShapeInfo,
|
|
void* vgradC0, const sd::LongType* gradC0ShapeInfo),
|
|
SD_FLOAT_TYPES);
|
|
|
|
//////////////////////////////////////////////////////////////////////////
|
|
void sruBIBP(LaunchContext* context, NDArray* x, NDArray* w, NDArray* b, NDArray* c0,
|
|
NDArray* ct, NDArray* gradCt, NDArray* gradHt, NDArray* mask, NDArray* gradI,
|
|
NDArray* gradW, NDArray* gradB, NDArray* gradC0) {
|
|
// x = x * mask
|
|
std::vector<LongType> dims = {1, 2};
|
|
if (mask) x->applyBroadcast(broadcast::Multiply, &dims, mask, x); // apply mask
|
|
|
|
// U = x * w
|
|
NDArray wi = mmul(*x, *w); // U [time x bS x 6*K]
|
|
|
|
const int time = x->sizeAt(0);
|
|
const int bS = x->sizeAt(1);
|
|
const int K = x->sizeAt(2) / 2;
|
|
|
|
std::vector<sd::LongType> gradBiasShape = {bS, 4 * K};
|
|
std::vector<sd::LongType> gradWiShape = {time, bS, 6 * K};
|
|
NDArray gradBias(x->ordering(), gradBiasShape, x->dataType(), context);
|
|
NDArray gradWi(x->ordering(), gradWiShape, x->dataType(), context);
|
|
|
|
PointersManager manager(context, "sru_bi_bp");
|
|
|
|
const int threadsPerBlock = SD_MAX_NUM_THREADS / 4;
|
|
const int blocksPerGrid = (x->sizeAt(1) * x->sizeAt(2) + threadsPerBlock - 1) /
|
|
threadsPerBlock; // loop through last two dimensions of x array -> bS, 2*K
|
|
const int sharedMem = threadsPerBlock * sizeof(LongType) * x->rankOf() + 128;
|
|
dim3 sruBiBpDims = sruBiDims(x->sizeAt(1) + x->sizeAt(2),x->rankOf());
|
|
NDArray::prepareSpecialUse({gradI, &gradWi, &gradBias, gradC0}, {x, &wi, b, c0, ct, gradCt, gradHt, mask});
|
|
BUILD_SINGLE_SELECTOR(
|
|
x->dataType(), sruBIBPCudaLauncher,
|
|
(sruBiBpDims.y, sruBiBpDims.x,sruBiBpDims.z, context->getCudaStream(), x->specialBuffer(), x->specialShapeInfo(),
|
|
wi.specialBuffer(), wi.specialShapeInfo(), b->specialBuffer(), b->specialShapeInfo(), c0->specialBuffer(),
|
|
c0->specialShapeInfo(), mask ? mask->specialBuffer() : nullptr, mask ? mask->specialShapeInfo() : nullptr,
|
|
ct->specialBuffer(), ct->specialShapeInfo(), gradHt->specialBuffer(), gradHt->specialShapeInfo(),
|
|
gradCt->specialBuffer(), gradCt->specialShapeInfo(), gradI->specialBuffer(), gradI->specialShapeInfo(),
|
|
gradWi.specialBuffer(), gradWi.specialShapeInfo(), gradBias.specialBuffer(), gradBias.specialShapeInfo(),
|
|
gradC0->specialBuffer(), gradC0->specialShapeInfo()),
|
|
SD_FLOAT_TYPES);
|
|
NDArray::registerSpecialUse({gradI, &gradWi, &gradBias, gradC0}, {x, &wi, b, c0, ct, gradCt, gradHt, mask});
|
|
|
|
manager.synchronize();
|
|
|
|
|
|
std::vector<LongType> dims2 = {0};
|
|
// gradB
|
|
gradBias.reduceAlongDimension(reduce::Sum, gradB, &dims2); // [4*K]
|
|
|
|
// gradW
|
|
x->permutei({0, 2, 1}, false, false); // [time, bS, 2*K] -> [time, 2*K, bS]
|
|
MmulHelper::mmul(x, &gradWi, gradW, 1., 0.); // [time, 2*K, bS] x [time, bS , 6*K] = [time, 2*K, 6*K]
|
|
}
|
|
|
|
} // namespace helpers
|
|
} // namespace ops
|
|
} // namespace sd
|