// File: 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:
subq $0xb8, %rsp
movq %rdi, 0xb0(%rsp)
movq 0xb0(%rsp), %rdi
movq %rdi, 0x8(%rsp)
movq (%rdi), %rax
callq *0x250(%rax)
cmpl $0x1, %eax
je 0xcb85
movq 0x8(%rsp), %rdi
movq (%rdi), %rax
callq *0x250(%rax)
cmpl $0x0, %eax
je 0xcb85
movq 0x8(%rsp), %rdi
movq (%rdi), %rax
callq *0x250(%rax)
cmpl $0x4, %eax
je 0xcb85
jmp 0xcd09
movq 0x8(%rsp), %rdi
callq 0xcd30
movq 0x8(%rsp), %rax
movq 0x400(%rax), %rdi
callq 0x1f310
movl %eax, 0xac(%rsp)
leaq 0x90(%rsp), %rdi
callq 0x1fac0
movl $0x0, 0x5c(%rsp)
cmpl $0x8, 0x5c(%rsp)
jge 0xcbef
movl 0xac(%rsp), %eax
movl $0x8, %ecx
cltd
idivl %ecx
movl %eax, %ecx
imull 0x5c(%rsp), %ecx
movslq 0x5c(%rsp), %rax
movl %ecx, 0x60(%rsp,%rax,4)
movl 0x5c(%rsp), %eax
addl $0x1, %eax
movl %eax, 0x5c(%rsp)
jmp 0xcbbc
movq 0x8(%rsp), %rax
movl 0xac(%rsp), %ecx
movl %ecx, 0x80(%rsp)
movq %rax, 0x50(%rsp)
movq $0x0, 0x48(%rsp)
cmpq $0x8, 0x48(%rsp)
jae 0xcc93
movq 0x48(%rsp), %rax
leaq 0x60(%rsp,%rax,4), %rdx
leaq 0x64(%rsp,%rax,4), %rcx
leaq 0x40(%rsp), %rdi
leaq 0x50(%rsp), %rsi
callq 0xce10
jmp 0xcc38
leaq 0x90(%rsp), %rdi
leaq 0x40(%rsp), %rsi
callq 0x1fad0
jmp 0xcc4c
leaq 0x40(%rsp), %rdi
callq 0x1fb00
movq 0x48(%rsp), %rax
addq $0x1, %rax
movq %rax, 0x48(%rsp)
jmp 0xcc10
movq %rax, %rcx
movl %edx, %eax
movq %rcx, 0x38(%rsp)
movl %eax, 0x34(%rsp)
jmp 0xcd11
movq %rax, %rcx
movl %edx, %eax
movq %rcx, 0x38(%rsp)
movl %eax, 0x34(%rsp)
leaq 0x40(%rsp), %rdi
callq 0x1fb00
jmp 0xcd11
leaq 0x90(%rsp), %rax
movq %rax, 0x28(%rsp)
movq 0x28(%rsp), %rdi
callq 0x1fb20
movq %rax, 0x20(%rsp)
movq 0x28(%rsp), %rdi
callq 0x1fb50
movq %rax, 0x18(%rsp)
leaq 0x20(%rsp), %rdi
leaq 0x18(%rsp), %rsi
callq 0x1fb80
testb $0x1, %al
jne 0xccd3
jmp 0xccfc
leaq 0x20(%rsp), %rdi
callq 0x1fbc0
movq %rax, 0x10(%rsp)
movq 0x10(%rsp), %rdi
callq 0x43c0
jmp 0xccee
jmp 0xccf0
leaq 0x20(%rsp), %rdi
callq 0x1fbd0
jmp 0xccbe
leaq 0x90(%rsp), %rdi
callq 0x1fc10
addq $0xb8, %rsp
retq
leaq 0x90(%rsp), %rdi
callq 0x1fc10
movq 0x38(%rsp), %rdi
callq 0x43f0
nopl (%rax,%rax)