Method exec_internal

dnn/src/common/multi_head_attn/proxy_backward_base.cpp:520–790

template <typename T>
void MHABackwardProxyBase::exec_internal(MHA_PROXY_BACKWARD_EXEC_PARAM) {
    auto wksp_bundle = get_workspace_bundle(
            MHA_PROXY_BACKWARD_TENSOR_TO_LAYOUT_CALL, workspace.raw_ptr);
    auto mask_bundle = get_mask_reservespace_bundle(
            MHA_PROXY_BACKWARD_TENSOR_TO_LAYOUT_CALL, mask_reservespace.raw_ptr());
    auto othr_bundle = get_othr_reservespace_bundle(
            MHA_PROXY_BACKWARD_TENSOR_TO_LAYOUT_CALL, othr_reservespace.raw_ptr());
    size_t head = param.num_heads;
    size_t one = 1;
    TensorND mask1{mask_bundle.get_workspace(0).raw_ptr, m_mask1_layout};
    TensorND mask2{mask_bundle.get_workspace(1).raw_ptr, m_mask2_layout};
    TensorND nq, nk, nv;
    if (param.qproj_size == 0 and param.num_heads == 1) {
        nq = queries;
    } else {
        nq = TensorND{othr_bundle.get_workspace(0).raw_ptr, m_grad_nq_layout};
    }
    if (param.kproj_size == 0 and param.num_heads == 1) {
        nk = keys;
    } else {
        nk = TensorND{othr_bundle.get_workspace(1).raw_ptr, m_grad_nk_layout};
    }
    if (param.vproj_size == 0 and param.num_heads == 1) {
        nv = values;
    } else {
        nv = TensorND{othr_bundle.get_workspace(2).raw_ptr, m_grad_nv_layout};
    }
    TensorND nx{othr_bundle.get_workspace(3).raw_ptr, m_grad_nx_layout};

    // out = dropout(out)
    TensorND grad_drop2{wksp_bundle.get_workspace(0).raw_ptr, m_grad_drop2_layout};
    m_dropoutbw_opr->param().drop_prob = param.out_prob;
    m_dropoutbw_opr->exec(diff, mask2, grad_drop2, wksp_bundle.get_workspace(1));

    // out = z @ wo + bo
    TensorND grad_z;
    if (param.oproj_size) {
        TensorND z{othr_bundle.get_workspace(4).raw_ptr, m_grad_z_layout};
        TensorND oweight{qkvo_weight_bias.ptr<T>() + m_wo_off, m_wo_layout};
        grad_z = TensorND{wksp_bundle.get_workspace(2).raw_ptr, m_grad_z_layout};
        TensorND grad_wo{wksp_bundle.get_workspace(3).raw_ptr, m_grad_wo_layout};
        m_matmul_opr->param().transposeA = false;
        m_matmul_opr->param().transposeB = true;
        matmul_exec(
                m_matmul_opr, grad_drop2, oweight, grad_z,
                wksp_bundle.get_workspace(4));
        m_bmatmul_opr->param().transposeA = true;
        m_bmatmul_opr->param().transposeB = false;
        m_bmatmul_opr->exec(z, grad_drop2, grad_wo, wksp_bundle.get_workspace(5));
        std::swap(m_grad_wo_layout.shape[0], one);
        TensorND doweight{dqkvo_weight_bias.ptr<T>() + m_wo_off, m_grad_wo_layout};
        std::swap(m_grad_wo_layout.shape[0], one);
        m_reduce_opr->param().axis = 0;
        m_reduce_opr->exec(grad_wo, doweight, wksp_bundle.get_workspace(6));
        if (param.obias) {
            TensorND dobias{dqkvo_weight_bias.ptr<T>() + m_bo_off, m_bo_layout};
            TensorND grad_bo{wksp_bundle.get_workspace(7).raw_ptr, m_grad_bo_layout};
            m_reduce_opr->exec(grad_drop2, grad_bo, wksp_bundle.get_workspace(8));