// File: /workspace/github_cpp_makefile_star2_O1/xbai0624[P]dnn/src/ConstructLayer.cpp
// Function: ConstructLayer::BatchNormalization_RestoreDelta()
void ConstructLayer::BatchNormalization_RestoreDelta()
{
    // generate the normalized delta: \hat{\delta}
    // refer to : https://arxiv.org/pdf/1502.03167v3.pdf 
    // for the detailed equation derivation, see the overleaf notes from Xinzhan Bai: xb4zp@virginia.edu

    // logic:
    // For layer n^th:
    // 1) if not use batch normalization, the \delta is back propagated directly from layer (n+1)^th
    //    \delta is defined as \partial{C}/\partial{z}
    // 2) if use batch normalization, the original (un-normalized) z does not participate in forward propagation,
    //    instead, the normalized z (or y) participated in forward propagation, 
    //    so the \delta back-propagated from (n+1)^th layer is actually \partial{C}/\partial{y} for n^th layer, 
    //    this is also true for the output layer : (the layer where the initial \delta was generated) 
    //    since the output layer used y to calculate cost and delta
    // 3) however, during backpropagation, take n^th -> (n+1)^th for example:
    //    because the y^n generated z^{n+1}, so we need the unormalized \partial{C}/\partial{z} in (n+1)^th layer to get
    //    \partial{C}/\partial{y} in n^th layer
    //
    // 4) so in compare with not-using-batch-normalization case, the using-batch-normalization case need one more step:
    //    derive \partial{C}/\partial{z} from \partial{C}/\partial{y}
    //    where \partial{C}/\partial{y} is back-propagated from (n+1)^th layer
    //    and \partial{C}/\partial{z} will be used for backpropagation to (n-1)^th layer
    //    This function is for the purpose of step 4)

    // only apply for cnn, fc and output layer, no need for pooling and input layers
    if( (GetType() != LayerType::cnn) &&
            ( GetType() != LayerType::fullyConnected) &&
            (GetType() != LayerType::output))
    {
        return; 
    }

    // extract gradients on batch normalization parameters
    // get ready for updating deltas
    BatchNormalization_GetGradientOnParameters();

    int batch_size =  __p_data_interface -> GetBatchSize();
#ifdef MULTI_THREAD
    std::vector<std::thread> vth;
    int Range[NTHREAD+1];
    for(int i=0;i<NTHREAD;i++)
        Range[i] = batch_size/NTHREAD * i;
    Range[NTHREAD] = batch_size;

    auto process_sample = [&](int start, int end)
    {
        for(int i=start;i<end;i++)
            BatchNormalization_RestoreDelta_ForSample(i);
    };

    for(size_t n=0;n<NTHREAD;n++)
        vth.push_back(std::thread(process_sample, Range[n], Range[n+1]));

    for(auto &th: vth)
        th.join();
#else
    for(int i=0;i<batch_size;i++)
        BatchNormalization_RestoreDelta_ForSample(i);
#endif
}

// Assembly:
pushq %rbp
pushq %r15
pushq %r14
pushq %r13
pushq %r12
pushq %rbx
subq $0x58, %rsp
movq %rdi, %rbx
movq (%rdi), %rax
callq *0x250(%rax)
cmpl $0x1, %eax
je 0x9064
movq (%rbx), %rax
movq %rbx, %rdi
callq *0x250(%rax)
testl %eax, %eax
je 0x9064
movq (%rbx), %rax
movq %rbx, %rdi
callq *0x250(%rax)
cmpl $0x4, %eax
jne 0x918f
movq %rbx, %rdi
callq 0x91e0
movq 0x400(%rbx), %rax
movl (%rax), %eax
xorps %xmm0, %xmm0
movaps %xmm0, 0x10(%rsp)
leal 0x7(%rax), %ecx
testl %eax, %eax
cmovnsl %eax, %ecx
movq $0x0, 0x20(%rsp)
movd %ecx, %xmm0
pshufd $0x0, %xmm0, %xmm0      # xmm0 = xmm0[0,0,0,0]
psrad $0x3, %xmm0
movdqa 0x1df7c(%rip), %xmm1    # 0x27020
xorl %ecx, %ecx
pshufd $0xf5, %xmm0, %xmm2     # xmm2 = xmm0[1,1,3,3]
movdqa 0x1df7d(%rip), %xmm3    # 0x27030
movdqa %xmm0, %xmm4
pmuludq %xmm1, %xmm4
pshufd $0xe8, %xmm4, %xmm4     # xmm4 = xmm4[0,2,2,3]
pshufd $0xf5, %xmm1, %xmm5     # xmm5 = xmm1[1,1,3,3]
pmuludq %xmm2, %xmm5
pshufd $0xe8, %xmm5, %xmm5     # xmm5 = xmm5[0,2,2,3]
punpckldq %xmm5, %xmm4    # xmm4 = xmm4[0],xmm5[0],xmm4[1],xmm5[1]
movdqa %xmm4, 0x30(%rsp,%rcx,4)
addq $0x4, %rcx
paddd %xmm3, %xmm1
cmpq $0x8, %rcx
jne 0x90b3
movl %eax, 0x50(%rsp)
xorl %r13d, %r13d
leaq 0x2d874(%rip), %rbp     # 0x36968
leaq 0x8(%rsp), %r14
movq %rsp, %r15
leaq 0x10(%rsp), %r12
movq $0x0, 0x8(%rsp)
movl $0x18, %edi
callq 0x4270
movq %rbp, (%rax)
movq 0x30(%rsp,%r13,4), %xmm0
pshufd $0xe1, %xmm0, %xmm0     # xmm0 = xmm0[1,0,2,3]
movq %xmm0, 0x8(%rax)
movq %rbx, 0x10(%rax)
movq %rax, (%rsp)
movq %r14, %rdi
movq %r15, %rsi
xorl %edx, %edx
callq 0x4210
movq (%rsp), %rdi
testq %rdi, %rdi
je 0x914c
movq (%rdi), %rax
callq *0x8(%rax)
movq %r12, %rdi
movq %r14, %rsi
callq 0x1726e
cmpq $0x0, 0x8(%rsp)
jne 0x91db
incq %r13
cmpq $0x8, %r13
jne 0x9101
movq 0x10(%rsp), %rbx
movq 0x18(%rsp), %r14
cmpq %r14, %rbx
je 0x9185
movq %rbx, %rdi
callq 0x43b0
addq $0x8, %rbx
jmp 0x9172
leaq 0x10(%rsp), %rdi
callq 0x13a40
addq $0x58, %rsp
popq %rbx
popq %r12
popq %r13
popq %r14
popq %r15
popq %rbp
retq
jmp 0x91c6
jmp 0x91c6
movq %rax, %rbx
movq (%rsp), %rdi
testq %rdi, %rdi
je 0x91b4
movq (%rdi), %rax
callq *0x8(%rax)
movq $0x0, (%rsp)
jmp 0x91c9
cmpq $0x0, 0x8(%rsp)
jne 0x91db
movq %rax, %rbx
leaq 0x10(%rsp), %rdi
callq 0x13a40
movq %rbx, %rdi
callq 0x43e0
callq 0x40a0