Function exec

dnn/src/common/rnn_cell.cpp:83–121

}

void exec(
        _megdnn_tensor_in input, _megdnn_tensor_in weight_ih, _megdnn_tensor_in bias_ih,
        _megdnn_tensor_in hx, _megdnn_tensor_in weight_hh, _megdnn_tensor_in bias_hh,
        _megdnn_tensor_out dst, _megdnn_workspace workspace,
        param::RNNCell::NonlineMode nonline_mode, Handle* handle) {
    TensorND tmp{static_cast<void*>(workspace.raw_ptr), dst.layout};
    _megdnn_workspace new_workspace = {
            workspace.raw_ptr + dst.layout.span().dist_byte(),
            workspace.size - dst.layout.span().dist_byte()};
    auto opr = handle->create_operator<MatrixMulForward>();
    opr->param().transposeB = true;
    opr->exec(input, weight_ih, tmp, new_workspace);
    opr->exec(hx, weight_hh, dst, new_workspace);
    auto add_opr = handle->create_operator<ElemwiseForward>();
    add_opr->param().mode = Elemwise::Param::Mode::ADD;
    add_opr->exec({dst, tmp}, dst);
    add_opr->exec({dst, bias_ih}, dst);
    add_opr->exec({dst, bias_hh}, dst);

    // activation
    using NonlineMode = param::RNNCell::NonlineMode;

    switch (nonline_mode) {
#define cb(_mode)                                                    \
    case NonlineMode::_mode: {                                       \
        auto nonlinear = handle->create_operator<ElemwiseForward>(); \
        nonlinear->param().mode = Elemwise::Param::Mode::_mode;      \
        nonlinear->exec({dst}, dst);                                 \
        break;                                                       \
    }
        cb(RELU);
        cb(TANH);
#undef cb
        case NonlineMode::IDENTITY:
            break;
        default:
            megdnn_assert(false);
    }
}

}  // namespace rnn_cell
}  // namespace megdnn