// File: /workspace/github_cpp_makefile_star2_O1/xbai0624[P]dnn/src/ConstructLayer.cpp
// Function: ConstructLayer::BatchNormalization_UpdateZ_2DLayer()
void ConstructLayer::BatchNormalization_UpdateZ_2DLayer()
{
    // 1) get mu and sigma_square
    // part I) active
    Images mu, sigma_square;
    GetMu_BFor2DLayer(__imageZ, mu);
    GetSigmaSquareFor2DLayer(__imageZ, mu, sigma_square);
    // part II) full
    Images mu_full, sigma_square_full;
    GetMu_BFor2DLayer(__imageZFull, mu_full);
    GetSigmaSquareFor2DLayer(__imageZFull, mu_full, sigma_square_full);

    // 2) save mu and sigma square
    // part I) active
    __imageAverage_BN[0] = mu;
    __imageVariance_BN[0] = sigma_square;
    // part II) full
    __imageAverageFull_BN[0] = mu_full;
    __imageVarianceFull_BN[0] = sigma_square_full;

    // 3) get normalized \hat{x} and save it to __imageZ_BN
    size_t batch_size = __imageZ.size();
    double epsilon = 1e-8; // avoid divide by zero issue; a hard-coded parameter from algorithm
    size_t nKernel = __imageZ[0].OutputImageFromKernel.size();
    // for parallelization
    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;

    /*
    for(size_t k=0;k<nKernel;k++)
    {
        // part I) active mu and sigma matrix
        Matrix & average_m_active = mu.OutputImageFromKernel[k];
        Matrix variance_m_active = sigma_square.OutputImageFromKernel[k];
        variance_m_active = variance_m_active + epsilon;
        variance_m_active(sqrt);
        // part II) full mu and sigma matrix
        Matrix & average_m_full = mu_full.OutputImageFromKernel[k];
        Matrix variance_m_full = sigma_square_full.OutputImageFromKernel[k];
        variance_m_full = variance_m_full + epsilon;
        variance_m_full(sqrt);

        // since one kernel corresponds to one pair of (\gamma, \beta)
        // so matrix average_m_active... etc must have dimension(1, 1)
        assert(average_m_active.Dimension().first == 1);
        assert(average_m_active.Dimension().second == 1);
        double _fMu_active = average_m_active[0][0];
        double _fVariance_active = variance_m_active[0][0];
        double _fMu_full = average_m_full[0][0];
        double _fVariance_full = variance_m_full[0][0];

        for(size_t i=0;i<batch_size;i++) // TO-DO: need to parallize
        {
            // part I) for active image
            Matrix m = __imageZ[i].OutputImageFromKernel[k]; // using copy assignment
            m = m - _fMu_active;
            m = m / _fVariance_active;
            __imageZ_BN[i].OutputImageFromKernel.push_back(m);

            // part II) for full image
            Matrix m_full = __imageZFull[i].OutputImageFromKernel[k];
            m_full = m_full - _fMu_full;
            m_full = m_full / _fVariance_full;
            __imageZFull_BN[i].OutputImageFromKernel.push_back(m_full);
        }
    }
    */

    // parallelize
    auto fill_z = [&](size_t start, size_t end)
    {
        for(size_t k=0;k<nKernel;k++)
        {
            // part I) active mu and sigma matrix
            Matrix & average_m_active = mu.OutputImageFromKernel[k];
            Matrix variance_m_active = sigma_square.OutputImageFromKernel[k];
            variance_m_active = variance_m_active + epsilon;
            variance_m_active(sqrt);
            // part II) full mu and sigma matrix
            Matrix & average_m_full = mu_full.OutputImageFromKernel[k];
            Matrix variance_m_full = sigma_square_full.OutputImageFromKernel[k];
            variance_m_full = variance_m_full + epsilon;
            variance_m_full(sqrt);

            // since one kernel corresponds to one pair of (\gamma, \beta)
            // so matrix average_m_active... etc must have dimension(1, 1)
            assert(average_m_active.Dimension().first == 1);
            assert(average_m_active.Dimension().second == 1);
            double _fMu_active = average_m_active[0][0];
            double _fVariance_active = variance_m_active[0][0];
            double _fMu_full = average_m_full[0][0];
            double _fVariance_full = variance_m_full[0][0];

            for(size_t i=start;i<end;i++) // TO-DO: need to parallize
            {
                // part I) for active image
                Matrix m = __imageZ[i].OutputImageFromKernel[k]; // using copy assignment
                m = m - _fMu_active;
                m = m / _fVariance_active;
                __imageZ_BN[i].OutputImageFromKernel.push_back(m);

                // part II) for full image
                Matrix m_full = __imageZFull[i].OutputImageFromKernel[k];
                m_full = m_full - _fMu_full;
                m_full = m_full / _fVariance_full;
                __imageZFull_BN[i].OutputImageFromKernel.push_back(m_full);
            }
        }
    };

#ifdef MULTI_THREAD
    vth.clear();
    for(int n = 0;n<NTHREAD;n++)
        vth.push_back(std::thread(fill_z, Range[n], Range[n+1]));
    for(auto &th: vth)
        th.join();
#else
    fill_z(0, batch_size);
#endif

    // 4) affine transformation to get y
    assert(__gamma_BN.size() == nKernel);
    assert(__beta_BN.size() == nKernel);

    /*
    for(size_t i = 0; i<batch_size;i++) // To-Do: need to parallize
    {
        for(size_t k=0;k<nKernel;k++)
        {
            double gamma = (__gamma_BN[k])[0][0];
            double beta = (__beta_BN[k])[0][0];
            // part I) active 
            Matrix m = __imageZ_BN[i].OutputImageFromKernel[k] * gamma;
            m = m + beta;
            __imageY_BN[i].OutputImageFromKernel.push_back(m);

            // part II) full
            m = __imageZFull_BN[i].OutputImageFromKernel[k] * gamma;
            m = m + beta;
            __imageYFull_BN[i].OutputImageFromKernel.push_back(m);
        }
    }
    */

    // parallelize
    auto fill_y = [&](size_t start, size_t end)
    {
        for(size_t i = start; i<end;i++) // To-Do: need to parallize
        {
            for(size_t k=0;k<nKernel;k++)
            {
                double gamma = (__gamma_BN[k])[0][0];
                double beta = (__beta_BN[k])[0][0];
                // part I) active 
                Matrix m = __imageZ_BN[i].OutputImageFromKernel[k] * gamma;
                m = m + beta;
                __imageY_BN[i].OutputImageFromKernel.push_back(m);

                // part II) full
                m = __imageZFull_BN[i].OutputImageFromKernel[k] * gamma;
                m = m + beta;
                __imageYFull_BN[i].OutputImageFromKernel.push_back(m);
            }
        }
    };
    vth.clear();
    for(int i=0;i<NTHREAD;i++)
        vth.push_back(std::thread(fill_y, Range[i], Range[i+1]));
    for(auto &i: vth) i.join();
}

// Assembly:
pushq %rbp
pushq %r15
pushq %r14
pushq %r13
pushq %r12
pushq %rbx
subq $0xd8, %rsp
movq %rdi, %rbx
pxor %xmm0, %xmm0
leaq 0x90(%rsp), %rdx
movdqa %xmm0, (%rdx)
xorl %eax, %eax
movq %rax, 0x10(%rdx)
movq %rax, 0x60(%rsp)
movdqa %xmm0, 0x50(%rsp)
leaq 0x150(%rdi), %r14
movq (%rdi), %rax
movq %r14, %rsi
callq *0x1b0(%rax)
movq (%rbx), %rax
leaq 0x90(%rsp), %rdx
leaq 0x50(%rsp), %rcx
movq %rbx, %rdi
movq %r14, %rsi
callq *0x1c0(%rax)
pxor %xmm0, %xmm0
leaq 0x70(%rsp), %rdx
movdqa %xmm0, (%rdx)
xorl %eax, %eax
movq %rax, 0x10(%rdx)
movq %rax, 0x40(%rsp)
movdqa %xmm0, 0x30(%rsp)
leaq 0x1b0(%rbx), %r14
movq (%rbx), %rax
movq %rbx, %rdi
movq %r14, %rsi
callq *0x1b0(%rax)
movq (%rbx), %rax
leaq 0x70(%rsp), %rdx
leaq 0x30(%rsp), %rcx
movq %rbx, %rdi
movq %r14, %rsi
callq *0x1c0(%rax)
movq 0x278(%rbx), %rdi
leaq 0x90(%rsp), %rsi
callq 0x1429e
movq 0x2a8(%rbx), %rdi
leaq 0x50(%rsp), %rsi
callq 0x1429e
movq 0x290(%rbx), %rdi
leaq 0x70(%rsp), %rsi
callq 0x1429e
movq 0x2c0(%rbx), %rdi
leaq 0x30(%rsp), %rsi
callq 0x1429e
movq 0x150(%rbx), %rcx
movq 0x158(%rbx), %rax
subq %rcx, %rax
sarq $0x3, %rax
movabsq $-0x5555555555555555, %rsi # imm = 0xAAAAAAAAAAAAAAAB
imulq %rsi, %rax
movabsq $0x3e45798ee2308c3a, %rdx # imm = 0x3E45798EE2308C3A
movq %rdx, 0xa8(%rsp)
movq 0x8(%rcx), %rdx
subq (%rcx), %rdx
sarq $0x3, %rdx
imulq %rsi, %rdx
movq %rdx, 0x28(%rsp)
pxor %xmm0, %xmm0
movdqa %xmm0, 0x10(%rsp)
movq $0x0, 0x20(%rsp)
movq %rax, %rcx
shrq $0x3, %rcx
movd %ecx, %xmm0
pshufd $0x0, %xmm0, %xmm0      # xmm0 = xmm0[0,0,0,0]
movdqa 0x1aa02(%rip), %xmm1    # 0x27020
xorl %ecx, %ecx
pshufd $0xf5, %xmm0, %xmm2     # xmm2 = xmm0[1,1,3,3]
movdqa 0x1aa03(%rip), %xmm3    # 0x27030
movdqa %xmm1, %xmm4
pmuludq %xmm0, %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, 0xb0(%rsp,%rcx,4)
addq $0x4, %rcx
paddd %xmm3, %xmm1
cmpq $0x8, %rcx
jne 0xc62d
movl %eax, 0xd0(%rsp)
movq 0x10(%rsp), %rax
movq 0x18(%rsp), %rcx
cmpq %rax, %rcx
je 0xc694
movq %rax, %rdx
cmpq $0x0, (%rdx)
jne 0xc99a
addq $0x8, %rdx
cmpq %rcx, %rdx
jne 0xc67c
movq %rax, 0x18(%rsp)
xorl %r12d, %r12d
leaq 0xa8(%rsp), %r13
leaq 0x8(%rsp), %rbp
movq %rsp, %r14
leaq 0x10(%rsp), %r15
movq $0x0, 0x8(%rsp)
movl $0x48, %edi
callq 0x4270
leaq 0x2a362(%rip), %rcx     # 0x36a28
movq %rcx, (%rax)
movq 0xb0(%rsp,%r12,4), %xmm0
pshufd $0xe1, %xmm0, %xmm0     # xmm0 = xmm0[1,0,2,3]
movq %xmm0, 0x8(%rax)
leaq 0x28(%rsp), %rcx
movq %rcx, 0x10(%rax)
leaq 0x90(%rsp), %rcx
movq %rcx, 0x18(%rax)
leaq 0x50(%rsp), %rcx
movq %rcx, 0x20(%rax)
movq %r13, 0x28(%rax)
leaq 0x70(%rsp), %rcx
movq %rcx, 0x30(%rax)
leaq 0x30(%rsp), %rcx
movq %rcx, 0x38(%rax)
movq %rbx, 0x40(%rax)
movq %rax, (%rsp)
movq %rbp, %rdi
movq %r14, %rsi
xorl %edx, %edx
callq 0x4210
movq (%rsp), %rdi
testq %rdi, %rdi
je 0xc735
movq (%rdi), %rax
callq *0x8(%rax)
movq %r15, %rdi
movq %rbp, %rsi
callq 0x1726e
cmpq $0x0, 0x8(%rsp)
jne 0xc99a
incq %r12
cmpq $0x8, %r12
jne 0xc6ac
movq 0x10(%rsp), %r14
movq 0x18(%rsp), %r15
cmpq %r15, %r14
movabsq $-0x5555555555555555, %r12 # imm = 0xAAAAAAAAAAAAAAAB
je 0xc783
movq %r14, %rdi
callq 0x43b0
addq $0x8, %r14
cmpq %r15, %r14
jne 0xc772
movq 0x310(%rbx), %rcx
subq 0x308(%rbx), %rcx
sarq $0x3, %rcx
imulq %r12, %rcx
movq 0x28(%rsp), %rax
cmpq %rax, %rcx
jne 0xc8e2
movq 0x328(%rbx), %rcx
subq 0x320(%rbx), %rcx
sarq $0x3, %rcx
imulq %r12, %rcx
cmpq %rax, %rcx
jne 0xc901
movq 0x10(%rsp), %rax
movq 0x18(%rsp), %rcx
cmpq %rax, %rcx
je 0xc7f0
movq %rax, %rdx
cmpq $0x0, (%rdx)
jne 0xc99a
addq $0x8, %rdx
cmpq %rcx, %rdx
jne 0xc7d8
movq %rax, 0x18(%rsp)
xorl %r13d, %r13d
leaq 0x2a26e(%rip), %rbp     # 0x36a68
leaq 0x8(%rsp), %r14
movq %rsp, %r15
leaq 0x10(%rsp), %r12
movq $0x0, 0x8(%rsp)
movl $0x20, %edi
callq 0x4270
movq %rbp, (%rax)
movq 0xb0(%rsp,%r13,4), %xmm0
pshufd $0xe1, %xmm0, %xmm0     # xmm0 = xmm0[1,0,2,3]
movq %xmm0, 0x8(%rax)
leaq 0x28(%rsp), %rcx
movq %rcx, 0x10(%rax)
movq %rbx, 0x18(%rax)
movq %rax, (%rsp)
movq %r14, %rdi
movq %r15, %rsi
xorl %edx, %edx
callq 0x4210
movq (%rsp), %rdi
testq %rdi, %rdi
je 0xc85e
movq (%rdi), %rax
callq *0x8(%rax)
movq %r12, %rdi
movq %r14, %rsi
callq 0x1726e
cmpq $0x0, 0x8(%rsp)
jne 0xc99a
incq %r13
cmpq $0x8, %r13
jne 0xc807
movq 0x10(%rsp), %rbx
movq 0x18(%rsp), %r14
cmpq %r14, %rbx
je 0xc89b
movq %rbx, %rdi
callq 0x43b0
addq $0x8, %rbx
jmp 0xc888
leaq 0x10(%rsp), %rdi
callq 0x13a40
leaq 0x30(%rsp), %rdi
callq 0x136aa
leaq 0x70(%rsp), %rdi
callq 0x136aa
leaq 0x50(%rsp), %rdi
callq 0x136aa
leaq 0x90(%rsp), %rdi
callq 0x136aa
addq $0xd8, %rsp
popq %rbx
popq %r12
popq %r13
popq %r14
popq %r15
popq %rbp
retq
leaq 0x1b87e(%rip), %rdi     # 0x28167
leaq 0x1ae5a(%rip), %rsi     # 0x2774a
leaq 0x1b88d(%rip), %rcx     # 0x28184
movl $0x8fe, %edx            # imm = 0x8FE
callq 0x4160
leaq 0x1b8be(%rip), %rdi     # 0x281c6
leaq 0x1ae3b(%rip), %rsi     # 0x2774a
leaq 0x1b86e(%rip), %rcx     # 0x28184
movl $0x8ff, %edx            # imm = 0x8FF
callq 0x4160
movq %rax, %rbx
jmp 0xc97b
movq %rax, %rbx
jmp 0xc967
jmp 0xc95a
jmp 0xc95a
jmp 0xc95a
jmp 0xc936
jmp 0xc952
jmp 0xc95a
movq %rax, %rbx
movq (%rsp), %rdi
testq %rdi, %rdi
je 0xc948
movq (%rdi), %rax
callq *0x8(%rax)
movq $0x0, (%rsp)
jmp 0xc95d
cmpq $0x0, 0x8(%rsp)
jne 0xc99a
movq %rax, %rbx
leaq 0x10(%rsp), %rdi
callq 0x13a40
leaq 0x30(%rsp), %rdi
callq 0x136aa
leaq 0x70(%rsp), %rdi
callq 0x136aa
leaq 0x50(%rsp), %rdi
callq 0x136aa
leaq 0x90(%rsp), %rdi
callq 0x136aa
movq %rbx, %rdi
callq 0x43e0
callq 0x40a0
nop