// File: yuang-chen[P]paraGCN/src/module.cpp
// Function: CrossEntropyLoss::forward(bool)
void CrossEntropyLoss::forward(bool training) {
    float total_loss = 0;
    int count = 0;
    if(training) {
        logits->zero_grad();
    }

    for(int i = 0; i < logits->data.size() / num_classes; i++) {
        if(truth[i] < 0) {
            continue;
        }
        count++;
        float *logit = &logits->data[i * num_classes];
        float max_logit = -1e30, sum_exp = 0;

        for(int j = 0; j < num_classes; j++) {
            max_logit = fmax(max_logit, logit[j]);
        }
        for(int j = 0; j < num_classes; j++) {
            logit[j] -= max_logit;
            sum_exp += expf(logit[j]);
        }
        total_loss += logf(sum_exp) - logit[truth[i]];

        if(training) {
            for(int j = 0; j < num_classes; j++) {
                float prob = expf(logit[j]) / sum_exp;
                logits->grad[i * num_classes + j] = prob;
            }
            logits->grad[i * num_classes + truth[i]] -= 1.0;
        }
    }
    *loss = total_loss / count;
    if(training) {
        for(int i = 0; i < logits->grad.size(); i++) {
            logits->grad[i] /= count;
        }
    }
}

// Assembly:
pushq %rbp
movq %rsp, %rbp
subq $0x70, %rsp
movb %sil, %al
movq %rdi, -0x8(%rbp)
andb $0x1, %al
movb %al, -0x9(%rbp)
movq -0x8(%rbp), %rax
movq %rax, -0x48(%rbp)
xorps %xmm0, %xmm0
movss %xmm0, -0x10(%rbp)
movl $0x0, -0x14(%rbp)
testb $0x1, -0x9(%rbp)
je 0xc2ae
movq -0x48(%rbp), %rax
movq 0x8(%rax), %rdi
callq 0xcff0
movl $0x0, -0x18(%rbp)
movq -0x48(%rbp), %rax
movslq -0x18(%rbp), %rcx
movq %rcx, -0x50(%rbp)
movq 0x8(%rax), %rdi
callq 0x5500
movq -0x48(%rbp), %rcx
movslq 0x20(%rcx), %rcx
xorl %edx, %edx
divq %rcx
movq %rax, %rcx
movq -0x50(%rbp), %rax
cmpq %rcx, %rax
jae 0xc4ba
movq -0x48(%rbp), %rax
movq 0x10(%rax), %rax
movslq -0x18(%rbp), %rcx
cmpl $0x0, (%rax,%rcx,4)
jge 0xc2fe
jmp 0xc4ac
movq -0x48(%rbp), %rcx
movl -0x14(%rbp), %eax
addl $0x1, %eax
movl %eax, -0x14(%rbp)
movq 0x8(%rcx), %rdi
movl -0x18(%rbp), %eax
imull 0x20(%rcx), %eax
movslq %eax, %rsi
callq 0x5520
movq %rax, -0x20(%rbp)
movss 0x3e9e(%rip), %xmm0     # 0x101c8
movss %xmm0, -0x24(%rbp)
xorps %xmm0, %xmm0
movss %xmm0, -0x28(%rbp)
movl $0x0, -0x2c(%rbp)
movq -0x48(%rbp), %rcx
movl -0x2c(%rbp), %eax
cmpl 0x20(%rcx), %eax
jge 0xc371
movss -0x24(%rbp), %xmm0
movq -0x20(%rbp), %rax
movslq -0x2c(%rbp), %rcx
movss (%rax,%rcx,4), %xmm1
callq 0xcb00
movss %xmm0, -0x24(%rbp)
movl -0x2c(%rbp), %eax
addl $0x1, %eax
movl %eax, -0x2c(%rbp)
jmp 0xc33e
movl $0x0, -0x30(%rbp)
movq -0x48(%rbp), %rcx
movl -0x30(%rbp), %eax
cmpl 0x20(%rcx), %eax
jge 0xc3c6
movss -0x24(%rbp), %xmm1
movq -0x20(%rbp), %rax
movslq -0x30(%rbp), %rcx
movss (%rax,%rcx,4), %xmm0
subss %xmm1, %xmm0
movss %xmm0, (%rax,%rcx,4)
movq -0x20(%rbp), %rax
movslq -0x30(%rbp), %rcx
movss (%rax,%rcx,4), %xmm0
callq 0x33e0
addss -0x28(%rbp), %xmm0
movss %xmm0, -0x28(%rbp)
movl -0x30(%rbp), %eax
addl $0x1, %eax
movl %eax, -0x30(%rbp)
jmp 0xc378
movss -0x28(%rbp), %xmm0
callq 0x3470
movq -0x48(%rbp), %rcx
movq -0x20(%rbp), %rax
movq 0x10(%rcx), %rcx
movslq -0x18(%rbp), %rdx
movslq (%rcx,%rdx,4), %rcx
subss (%rax,%rcx,4), %xmm0
addss -0x10(%rbp), %xmm0
movss %xmm0, -0x10(%rbp)
testb $0x1, -0x9(%rbp)
je 0xc4aa
movl $0x0, -0x34(%rbp)
movq -0x48(%rbp), %rcx
movl -0x34(%rbp), %eax
cmpl 0x20(%rcx), %eax
jge 0xc468
movq -0x20(%rbp), %rax
movslq -0x34(%rbp), %rcx
movss (%rax,%rcx,4), %xmm0
callq 0x33e0
movq -0x48(%rbp), %rcx
divss -0x28(%rbp), %xmm0
movss %xmm0, -0x38(%rbp)
movss -0x38(%rbp), %xmm0
movss %xmm0, -0x54(%rbp)
movq 0x8(%rcx), %rdi
addq $0x18, %rdi
movl -0x18(%rbp), %eax
imull 0x20(%rcx), %eax
addl -0x34(%rbp), %eax
movslq %eax, %rsi
callq 0x5520
movss -0x54(%rbp), %xmm0
movss %xmm0, (%rax)
movl -0x34(%rbp), %eax
addl $0x1, %eax
movl %eax, -0x34(%rbp)
jmp 0xc404
movq -0x48(%rbp), %rcx
movq 0x8(%rcx), %rdi
addq $0x18, %rdi
movl -0x18(%rbp), %eax
imull 0x20(%rcx), %eax
movq 0x10(%rcx), %rcx
movslq -0x18(%rbp), %rdx
addl (%rcx,%rdx,4), %eax
movslq %eax, %rsi
callq 0x5520
movss (%rax), %xmm0
cvtss2sd %xmm0, %xmm0
movsd 0x3d22(%rip), %xmm1     # 0x101c0
subsd %xmm1, %xmm0
cvtsd2ss %xmm0, %xmm0
movss %xmm0, (%rax)
jmp 0xc4ac
movl -0x18(%rbp), %eax
addl $0x1, %eax
movl %eax, -0x18(%rbp)
jmp 0xc2b5
movq -0x48(%rbp), %rax
movss -0x10(%rbp), %xmm0
cvtsi2ssl -0x14(%rbp), %xmm1
divss %xmm1, %xmm0
movq 0x18(%rax), %rax
movss %xmm0, (%rax)
testb $0x1, -0x9(%rbp)
je 0xc543
movl $0x0, -0x3c(%rbp)
movq -0x48(%rbp), %rax
movslq -0x3c(%rbp), %rcx
movq %rcx, -0x60(%rbp)
movq 0x8(%rax), %rdi
addq $0x18, %rdi
callq 0x5500
movq %rax, %rcx
movq -0x60(%rbp), %rax
cmpq %rcx, %rax
jae 0xc541
movq -0x48(%rbp), %rax
cvtsi2ssl -0x14(%rbp), %xmm0
movss %xmm0, -0x64(%rbp)
movq 0x8(%rax), %rdi
addq $0x18, %rdi
movslq -0x3c(%rbp), %rsi
callq 0x5520
movss -0x64(%rbp), %xmm1
movss (%rax), %xmm0
divss %xmm1, %xmm0
movss %xmm0, (%rax)
movl -0x3c(%rbp), %eax
addl $0x1, %eax
movl %eax, -0x3c(%rbp)
jmp 0xc4e1
jmp 0xc543
addq $0x70, %rsp
popq %rbp
retq
nopl (%rax)