// File: /workspace/github_cpp_makefile_star2_O3/xbai0624[P]dnn/src/ConstructLayer.cpp
// Function: ConstructLayer::BatchNormalization_UpdateZ_1DLayer()
void ConstructLayer::BatchNormalization_UpdateZ_1DLayer()
{
    // 1) get mu and sigma_square
    // part I) active
    Images mu, sigma_square;
    GetMu_BFor1DLayer(__imageZ, mu); // use active z image
    GetSigmaSquareFor1DLayer(__imageZ, mu, sigma_square);
    // part II) full
    Images mu_full, sigma_square_full;
    GetMu_BFor1DLayer(__imageZFull, mu_full);
    GetSigmaSquareFor1DLayer(__imageZFull, mu_full, sigma_square_full);

    // 2) save mu and sigma_square ---- Do I need to calculate both active and full? This is time consuming.
    // 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();
    assert(__imageZ[0].OutputImageFromKernel.size() == 1); // make sure 1D layer
    double epsilon = 1e-8; // avoid divide by zero issue; a hard-coded parameter from algorithm

    // part I) active mu and sigma matrix
    Matrix & average_m_active = mu.OutputImageFromKernel[0];
    Matrix variance_m_active = sigma_square.OutputImageFromKernel[0];
    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[0];
    Matrix variance_m_full = sigma_square_full.OutputImageFromKernel[0];
    variance_m_full = variance_m_full + epsilon;
    variance_m_full(sqrt);
    /*
    for(size_t i=0;i<batch_size;i++) // TO-DO: need to parallize
    {
        // part I) for active image
        Matrix m = __imageZ[i].OutputImageFromKernel[0]; // using copy assigment
        m = m - average_m_active;
        m = m / variance_m_active;
        Images tmp;
        tmp.OutputImageFromKernel.push_back(m);
        __imageZ_BN[i]=tmp;


        // part II) for full image
        Matrix m_full = __imageZFull[i].OutputImageFromKernel[0]; // using copy assigment
        m_full = m_full - average_m_full;
        m_full = m_full / variance_m_full;
        Images tmp_full;
        tmp_full.OutputImageFromKernel.push_back(m_full);
        __imageZFull_BN[i]=tmp_full;
    }
    */
    // parallelize
    auto fill_z = [&](size_t start, size_t end) 
    {
        for(size_t i=start;i<end;i++)
        {
            // part I) for active image
            Matrix m = __imageZ[i].OutputImageFromKernel[0]; // using copy assigment
            m = m - average_m_active;
            m = m / variance_m_active;
            Images tmp;
            tmp.OutputImageFromKernel.push_back(m);
            __imageZ_BN[i]=tmp;


            // part II) for full image
            Matrix m_full = __imageZFull[i].OutputImageFromKernel[0]; // using copy assigment
            m_full = m_full - average_m_full;
            m_full = m_full / variance_m_full;
            Images tmp_full;
            tmp_full.OutputImageFromKernel.push_back(m_full);
            __imageZFull_BN[i]=tmp_full;
        }
    };

#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;

    for(size_t 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() == 1); // for 1D layer, size = 1; for 2D layer, size = nKernel
    assert(__beta_BN.size() == 1); 
    /*
    for(size_t i=0;i<batch_size;i++) // TO-DO: need to parallize
    {
        // part I) active
        Matrix m = __imageZ_BN[i].OutputImageFromKernel[0] ^ __gamma_BN[0];
        m = m + __beta_BN[0];
        Images tmp;
        tmp.OutputImageFromKernel.push_back(m);
        __imageY_BN[i]=tmp;

        // part II) full
        m = __imageZFull_BN[i].OutputImageFromKernel[0] ^ __gamma_BN[0];
        m = m + __beta_BN[0];
        tmp.OutputImageFromKernel.clear();
        tmp.OutputImageFromKernel.push_back(m);
        __imageYFull_BN[i]=tmp;
    }
    */
    // parallelize
    auto fill_y = [&](size_t start, size_t end)
    {
        for(size_t i=start;i<end;i++)
        {
            // part I) active
            Matrix m = __imageZ_BN[i].OutputImageFromKernel[0] ^ __gamma_BN[0];
            m = m + __beta_BN[0];
            Images tmp;
            tmp.OutputImageFromKernel.push_back(m);
            __imageY_BN[i]=tmp;

            // part II) full
            m = __imageZFull_BN[i].OutputImageFromKernel[0] ^ __gamma_BN[0];
            m = m + __beta_BN[0];
            tmp.OutputImageFromKernel.clear();
            tmp.OutputImageFromKernel.push_back(m);
            __imageYFull_BN[i]=tmp;
        }
    };

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

// Assembly:
pushq %rbp
pushq %r15
pushq %r14
pushq %r13
pushq %r12
pushq %rbx
subq $0x118, %rsp            # imm = 0x118
movq %rdi, %r15
pxor %xmm0, %xmm0
leaq 0xf0(%rsp), %rdx
movdqa %xmm0, (%rdx)
xorl %eax, %eax
movq %rax, 0x10(%rdx)
movq %rax, 0x90(%rsp)
movdqa %xmm0, 0x80(%rsp)
leaq 0x150(%rdi), %r14
movq (%rdi), %rax
movq %r14, %rsi
callq *0x1a8(%rax)
movq (%r15), %rax
leaq 0xf0(%rsp), %rdx
leaq 0x80(%rsp), %rcx
movq %r15, %rdi
movq %r14, %rsi
callq *0x1b8(%rax)
pxor %xmm0, %xmm0
leaq 0xd0(%rsp), %rdx
movdqa %xmm0, (%rdx)
xorl %eax, %eax
movq %rax, 0x10(%rdx)
movq %rax, 0x70(%rsp)
movdqa %xmm0, 0x60(%rsp)
leaq 0x1b0(%r15), %r14
movq (%r15), %rax
movq %r15, %rdi
movq %r14, %rsi
callq *0x1a8(%rax)
movq (%r15), %rax
leaq 0xd0(%rsp), %rdx
leaq 0x60(%rsp), %rcx
movq %r15, %rdi
movq %r14, %rsi
callq *0x1b8(%rax)
movq 0x278(%r15), %rdi
leaq 0xf0(%rsp), %rsi
callq 0x14112
movq 0x2a8(%r15), %rdi
leaq 0x80(%rsp), %rsi
callq 0x14112
movq 0x290(%r15), %rdi
leaq 0xd0(%rsp), %rsi
callq 0x14112
movq 0x2c0(%r15), %rdi
leaq 0x60(%rsp), %rsi
callq 0x14112
movq 0x150(%r15), %rbx
movq 0x8(%rbx), %rax
subq (%rbx), %rax
cmpq $0x18, %rax
jne 0xc338
movq 0x158(%r15), %r12
movq 0xf0(%rsp), %rax
movq %rax, 0x110(%rsp)
movq 0x80(%rsp), %rsi
leaq 0xb8(%rsp), %rdi
callq 0x1aa3e
leaq 0x30(%rsp), %rdi
leaq 0xb8(%rsp), %rsi
movq 0x1b5d6(%rip), %xmm0    # 0x27598
callq 0x1b3f2
leaq 0xb8(%rsp), %rdi
leaq 0x30(%rsp), %rsi
callq 0x14e2c
leaq 0x30(%rsp), %rdi
callq 0x1abf2
movq 0x2afd6(%rip), %rsi     # 0x36fc0
leaq 0xb8(%rsp), %rdi
callq 0x1bf18
movq 0xd0(%rsp), %rax
movq %rax, 0x108(%rsp)
movq 0x60(%rsp), %rsi
leaq 0xa0(%rsp), %rdi
callq 0x1aa3e
leaq 0x30(%rsp), %rdi
leaq 0xa0(%rsp), %rsi
movq 0x1b56a(%rip), %xmm0    # 0x27598
callq 0x1b3f2
leaq 0xa0(%rsp), %rdi
leaq 0x30(%rsp), %rsi
callq 0x14e2c
leaq 0x30(%rsp), %rdi
callq 0x1abf2
movq 0x2af6a(%rip), %rsi     # 0x36fc0
leaq 0xa0(%rsp), %rdi
callq 0x1bf18
movq %r15, 0x28(%rsp)
subq %rbx, %r12
sarq $0x3, %r12
movabsq $-0x5555555555555555, %rax # imm = 0xAAAAAAAAAAAAAAAB
imulq %r12, %rax
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 0x1af78(%rip), %xmm1    # 0x27020
xorl %ecx, %ecx
pshufd $0xf5, %xmm0, %xmm2     # xmm2 = xmm0[1,1,3,3]
movdqa 0x1af79(%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, 0x30(%rsp,%rcx,4)
addq $0x4, %rcx
paddd %xmm3, %xmm1
cmpq $0x8, %rcx
jne 0xc0b7
movl %eax, 0x50(%rsp)
movl $0x1, %r14d
movl 0x30(%rsp), %ebp
movq %rsp, %r12
leaq 0x10(%rsp), %rbx
movl %ebp, %r15d
movl 0x30(%rsp,%r14,4), %r13d
movq $0x0, (%rsp)
movl $0x38, %edi
callq 0x4270
leaq 0x2a887(%rip), %rcx     # 0x369a8
movq %rcx, (%rax)
movl %r13d, 0x8(%rax)
movl %r15d, 0xc(%rax)
movq 0x28(%rsp), %rcx
movq %rcx, 0x10(%rax)
movq 0x110(%rsp), %rcx
movq %rcx, 0x18(%rax)
leaq 0xb8(%rsp), %rcx
movq %rcx, 0x20(%rax)
movq 0x108(%rsp), %rcx
movq %rcx, 0x28(%rax)
leaq 0xa0(%rsp), %rcx
movq %rcx, 0x30(%rax)
movq %rax, 0x8(%rsp)
movq %r12, %rdi
leaq 0x8(%rsp), %rsi
xorl %edx, %edx
callq 0x4210
movq 0x8(%rsp), %rdi
testq %rdi, %rdi
je 0xc189
movq (%rdi), %rax
callq *0x8(%rax)
movq %rbx, %rdi
movq %r12, %rsi
callq 0x1718e
cmpq $0x0, (%rsp)
jne 0xc452
incq %r14
movl %r13d, %r15d
cmpq $0x9, %r14
jne 0xc103
movq 0x10(%rsp), %r14
movq 0x18(%rsp), %rbx
cmpq %rbx, %r14
movq 0x28(%rsp), %r15
je 0xc1d4
movq %r14, %rdi
callq 0x43b0
addq $0x8, %r14
cmpq %rbx, %r14
jne 0xc1c3
movq 0x310(%r15), %rax
subq 0x308(%r15), %rax
cmpq $0x18, %rax
jne 0xc357
movq 0x328(%r15), %rax
subq 0x320(%r15), %rax
cmpq $0x18, %rax
jne 0xc376
movq 0x10(%rsp), %rax
movq 0x18(%rsp), %rcx
cmpq %rax, %rcx
je 0xc22e
movq %rax, %rdx
cmpq $0x0, (%rdx)
jne 0xc452
addq $0x8, %rdx
cmpq %rcx, %rdx
jne 0xc216
movq %rax, 0x18(%rsp)
movl $0x1, %r13d
movq %rsp, %r14
leaq 0x8(%rsp), %r15
leaq 0x10(%rsp), %r12
movl 0x30(%rsp,%r13,4), %ebx
movq $0x0, (%rsp)
movl $0x18, %edi
callq 0x4270
leaq 0x2a789(%rip), %rcx     # 0x369e8
movq %rcx, (%rax)
movl %ebx, 0x8(%rax)
movl %ebp, 0xc(%rax)
movq 0x28(%rsp), %rcx
movq %rcx, 0x10(%rax)
movq %rax, 0x8(%rsp)
movq %r14, %rdi
movq %r15, %rsi
xorl %edx, %edx
callq 0x4210
movq 0x8(%rsp), %rdi
testq %rdi, %rdi
je 0xc293
movq (%rdi), %rax
callq *0x8(%rax)
movq %r12, %rdi
movq %r14, %rsi
callq 0x1718e
cmpq $0x0, (%rsp)
jne 0xc452
incq %r13
movl %ebx, %ebp
cmpq $0x9, %r13
jne 0xc241
movq 0x10(%rsp), %rbx
movq 0x18(%rsp), %r14
cmpq %r14, %rbx
je 0xc2d1
movq %rbx, %rdi
callq 0x43b0
addq $0x8, %rbx
jmp 0xc2be
leaq 0x10(%rsp), %rdi
callq 0x138d0
leaq 0xa0(%rsp), %rdi
callq 0x1abf2
leaq 0xb8(%rsp), %rdi
callq 0x1abf2
leaq 0x60(%rsp), %rdi
callq 0x1353a
leaq 0xd0(%rsp), %rdi
callq 0x1353a
leaq 0x80(%rsp), %rdi
callq 0x1353a
leaq 0xf0(%rsp), %rdi
callq 0x1353a
addq $0x118, %rsp            # imm = 0x118
popq %rbx
popq %r12
popq %r13
popq %r14
popq %r15
popq %rbp
retq
leaq 0x1bd4c(%rip), %rdi     # 0x2808b
leaq 0x1b404(%rip), %rsi     # 0x2774a
leaq 0x1bd6c(%rip), %rcx     # 0x280b9
movl $0x801, %edx            # imm = 0x801
callq 0x4160
leaq 0x1bd9d(%rip), %rdi     # 0x280fb
leaq 0x1b3e5(%rip), %rsi     # 0x2774a
leaq 0x1bd4d(%rip), %rcx     # 0x280b9
movl $0x84c, %edx            # imm = 0x84C
callq 0x4160
leaq 0x1bd95(%rip), %rdi     # 0x28112
leaq 0x1b3c6(%rip), %rsi     # 0x2774a
leaq 0x1bd2e(%rip), %rcx     # 0x280b9
movl $0x84d, %edx            # imm = 0x84D
callq 0x4160
jmp 0xc3a6
movq %rax, %rbx
leaq 0x30(%rsp), %rdi
callq 0x1abf2
jmp 0xc3ff
movq %rax, %rbx
jmp 0xc3ff
jmp 0xc3b9
jmp 0xc3b9
movq %rax, %rbx
leaq 0x30(%rsp), %rdi
jmp 0xc407
movq %rax, %rbx
jmp 0xc40c
jmp 0xc3c5
movq %rax, %rbx
jmp 0xc430
movq %rax, %rbx
jmp 0xc419
jmp 0xc3f2
jmp 0xc3f2
jmp 0xc3f2
jmp 0xc3d6
jmp 0xc3eb
jmp 0xc3f2
movq %rax, %rbx
movq 0x8(%rsp), %rdi
testq %rdi, %rdi
je 0xc3f5
movq (%rdi), %rax
callq *0x8(%rax)
jmp 0xc3f5
cmpq $0x0, (%rsp)
jne 0xc452
movq %rax, %rbx
leaq 0x10(%rsp), %rdi
callq 0x138d0
leaq 0xa0(%rsp), %rdi
callq 0x1abf2
leaq 0xb8(%rsp), %rdi
callq 0x1abf2
leaq 0x60(%rsp), %rdi
callq 0x1353a
leaq 0xd0(%rsp), %rdi
callq 0x1353a
leaq 0x80(%rsp), %rdi
callq 0x1353a
leaq 0xf0(%rsp), %rdi
callq 0x1353a
movq %rbx, %rdi
callq 0x43e0
callq 0x40a0
nop