// File: /workspace/github_cpp_makefile_star2_O2/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 0x85ac
movq (%rbx), %rax
movq %rbx, %rdi
callq *0x250(%rax)
testl %eax, %eax
je 0x85ac
movq (%rbx), %rax
movq %rbx, %rdi
callq *0x250(%rax)
cmpl $0x4, %eax
jne 0x868e
movq %rbx, %rdi
callq 0x86d8
movq 0x400(%rbx), %rax
movl (%rax), %ecx
pxor %xmm0, %xmm0
andq $0x0, 0x20(%rsp)
movdqa %xmm0, 0x10(%rsp)
pushq $0x8
popq %rsi
movl %ecx, %eax
cltd
idivl %esi
xorl %edx, %edx
xorl %esi, %esi
cmpq $0x8, %rsi
je 0x85ea
movl %edx, 0x30(%rsp,%rsi,4)
incq %rsi
addl %eax, %edx
jmp 0x85d9
movl %ecx, 0x50(%rsp)
xorl %ebp, %ebp
leaq 0x2a379(%rip), %r14     # 0x32970
leaq 0x8(%rsp), %r15
movq %rsp, %r12
leaq 0x10(%rsp), %r13
cmpq $0x8, %rbp
je 0x8667
andq $0x0, 0x8(%rsp)
pushq $0x18
popq %rdi
callq 0x4250
movq %r14, (%rax)
movq 0x30(%rsp,%rbp,4), %xmm0
pshufd $0xe1, %xmm0, %xmm0     # xmm0 = xmm0[1,0,2,3]
movq %xmm0, 0x8(%rax)
movq %rbx, 0x10(%rax)
movq %rax, (%rsp)
movq %r15, %rdi
movq %r12, %rsi
xorl %edx, %edx
callq 0x4200
movq (%rsp), %rdi
testq %rdi, %rdi
je 0x864f
movq (%rdi), %rax
callq *0x8(%rax)
movq %r13, %rdi
movq %r15, %rsi
callq 0x1470a
movq %r15, %rdi
callq 0x113c8
incq %rbp
jmp 0x8604
movq 0x10(%rsp), %rbx
movq 0x18(%rsp), %r14
cmpq %r14, %rbx
je 0x8684
movq %rbx, %rdi
callq 0x4390
addq $0x8, %rbx
jmp 0x8671
leaq 0x10(%rsp), %rdi
callq 0x113de
addq $0x58, %rsp
popq %rbx
popq %r12
popq %r13
popq %r14
popq %r15
popq %rbp
retq
jmp 0x869f
movq %rax, %rbx
jmp 0x86c5
movq %rax, %rbx
movq (%rsp), %rdi
testq %rdi, %rdi
je 0x86c5
movq (%rdi), %rax
callq *0x8(%rax)
jmp 0x86c5
movq %rax, %rbx
leaq 0x8(%rsp), %rdi
callq 0x113c8
leaq 0x10(%rsp), %rdi
callq 0x113de
movq %rbx, %rdi
callq 0x43b0
nop