// File: /workspace/github_cpp_makefile_star2_O2/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
pushq %r15
pushq %r14
pushq %r13
pushq %r12
pushq %rbx
subq $0x28, %rsp
movq %rdi, %r14
movl %esi, 0x18(%rsp)
testl %esi, %esi
je 0x6208
movq 0x8(%r14), %rdi
callq 0x673e
xorps %xmm0, %xmm0
movss %xmm0, 0x10(%rsp)
xorl %r15d, %r15d
movl $0x0, 0x14(%rsp)
movq 0x8(%r14), %rsi
movq (%rsi), %rcx
movq 0x8(%rsi), %rax
subq %rcx, %rax
sarq $0x2, %rax
movslq 0x20(%r14), %rbx
xorl %edx, %edx
divq %rbx
cmpq %r15, %rax
jbe 0x6391
movq 0x10(%r14), %rax
cmpl $0x0, (%rax,%r15,4)
js 0x6389
movl %ebx, %eax
imull %r15d, %eax
cltq
leaq (%rcx,%rax,4), %rax
movq %rax, 0x20(%rsp)
testl %ebx, %ebx
movl $0x0, %ebp
cmovgl %ebx, %ebp
movl %ebx, %eax
imull %r15d, %eax
cltq
leaq (%rcx,%rax,4), %r12
xorl %r13d, %r13d
movss 0x1f63(%rip), %xmm2     # 0x81e4
cmpq %r13, %rbp
je 0x629c
movss (%r12,%r13,4), %xmm1
movaps %xmm2, %xmm0
callq 0x3210
movaps %xmm0, %xmm2
incq %r13
jmp 0x6281
xorps %xmm0, %xmm0
movss %xmm0, 0xc(%rsp)
xorl %ebp, %ebp
movss %xmm2, 0x1c(%rsp)
movslq %ebx, %rax
cmpq %rax, %rbp
jge 0x62e9
movss (%r12,%rbp,4), %xmm0
subss %xmm2, %xmm0
movss %xmm0, (%r12,%rbp,4)
callq 0x3330
movss 0x1c(%rsp), %xmm2
movss 0xc(%rsp), %xmm1
addss %xmm0, %xmm1
movss %xmm1, 0xc(%rsp)
incq %rbp
movl 0x20(%r14), %ebx
jmp 0x62ad
incl 0x14(%rsp)
movss 0xc(%rsp), %xmm0
callq 0x3390
movq 0x10(%r14), %rbx
movslq (%rbx,%r15,4), %rax
movq 0x20(%rsp), %rcx
subss (%rcx,%rax,4), %xmm0
movss 0x10(%rsp), %xmm1
addss %xmm0, %xmm1
movss %xmm1, 0x10(%rsp)
cmpb $0x0, 0x18(%rsp)
je 0x6389
movq 0x8(%r14), %r13
movl 0x20(%r14), %eax
xorl %ebp, %ebp
movslq %eax, %rcx
cmpq %rcx, %rbp
jge 0x6362
movss (%r12,%rbp,4), %xmm0
callq 0x3330
divss 0xc(%rsp), %xmm0
movl 0x20(%r14), %eax
movl %r15d, %ecx
imull %eax, %ecx
movslq %ecx, %rcx
addq %rbp, %rcx
movq 0x18(%r13), %rdx
movss %xmm0, (%rdx,%rcx,4)
incq %rbp
jmp 0x632b
movq 0x8(%r14), %rcx
imull %r15d, %eax
movslq (%rbx,%r15,4), %rdx
cltq
addq %rdx, %rax
movq 0x18(%rcx), %rcx
movss (%rcx,%rax,4), %xmm0
addss 0x1e64(%rip), %xmm0     # 0x81e8
movss %xmm0, (%rcx,%rax,4)
incq %r15
jmp 0x621c
cvtsi2ssl 0x14(%rsp), %xmm0
movss 0x10(%rsp), %xmm1
divss %xmm0, %xmm1
movq 0x18(%r14), %rax
movss %xmm1, (%rax)
cmpb $0x0, 0x18(%rsp)
je 0x63d9
movq 0x18(%rsi), %rax
movq 0x20(%rsi), %rcx
subq %rax, %rcx
sarq $0x2, %rcx
xorl %edx, %edx
cmpq %rdx, %rcx
je 0x63d9
movss (%rax,%rdx,4), %xmm1
divss %xmm0, %xmm1
movss %xmm1, (%rax,%rdx,4)
incq %rdx
jmp 0x63c1
addq $0x28, %rsp
popq %rbx
popq %r12
popq %r13
popq %r14
popq %r15
popq %rbp
retq