// File: /workspace/github_cpp_makefile_star2_O3/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, 0xc(%rsp)
testl %esi, %esi
je 0x679c
movq 0x8(%r14), %rdi
callq 0x6d9c
movq 0x8(%r14), %r13
movq (%r13), %rbx
movq 0x8(%r13), %rax
subq %rbx, %rax
sarq $0x2, %rax
movslq 0x20(%r14), %r12
xorps %xmm0, %xmm0
cmpq %r12, %rax
jae 0x67c3
xorps %xmm1, %xmm1
jmp 0x697a
movq 0x10(%r14), %rsi
xorl %r15d, %r15d
xorps %xmm1, %xmm1
movl $0x0, 0x8(%rsp)
cmpl $0x0, (%rsi,%r15,4)
js 0x694e
movss %xmm1, (%rsp)
xorps %xmm3, %xmm3
testl %r12d, %r12d
jle 0x687d
movl %r12d, %eax
movl %r12d, %ecx
imull %r15d, %ecx
movslq %ecx, %rcx
leaq (%rbx,%rcx,4), %r13
xorl %ecx, %ecx
movss 0x19d8(%rip), %xmm0     # 0x81e4
movaps %xmm0, %xmm4
movss (%r13,%rcx,4), %xmm1
cmpunordss %xmm4, %xmm4
movaps %xmm4, %xmm2
andps %xmm1, %xmm2
maxss %xmm0, %xmm1
andnps %xmm1, %xmm4
orps %xmm2, %xmm4
incq %rcx
movaps %xmm4, %xmm0
cmpq %rcx, %rax
jne 0x680f
testl %r12d, %r12d
jle 0x687d
xorps %xmm3, %xmm3
xorl %ebp, %ebp
movaps %xmm4, 0x10(%rsp)
movss %xmm3, 0x4(%rsp)
movss (%r13,%rbp,4), %xmm0
subss %xmm4, %xmm0
movss %xmm0, (%r13,%rbp,4)
callq 0x3330
movaps 0x10(%rsp), %xmm4
movss 0x4(%rsp), %xmm3
addss %xmm0, %xmm3
incq %rbp
movslq 0x20(%r14), %rax
cmpq %rax, %rbp
jl 0x6840
incl 0x8(%rsp)
movl %r12d, %eax
imull %r15d, %eax
cltq
leaq (%rbx,%rax,4), %r13
movss %xmm3, 0x4(%rsp)
movaps %xmm3, %xmm0
callq 0x33a0
movq 0x10(%r14), %rsi
movslq (%rsi,%r15,4), %rcx
subss (%r13,%rcx,4), %xmm0
movq 0x8(%r14), %r13
movss (%rsp), %xmm1
addss %xmm0, %xmm1
movl 0x20(%r14), %eax
cmpb $0x0, 0xc(%rsp)
je 0x694b
movq 0x18(%r13), %rbp
testl %eax, %eax
jle 0x692c
movq %rsi, 0x10(%rsp)
movss %xmm1, (%rsp)
imull %r15d, %r12d
movslq %r12d, %rax
leaq (%rbx,%rax,4), %rbx
xorl %r12d, %r12d
movss (%rbx,%r12,4), %xmm0
callq 0x3330
divss 0x4(%rsp), %xmm0
movslq 0x20(%r14), %rax
movl %r15d, %ecx
imull %eax, %ecx
movslq %ecx, %rcx
addq %r12, %rcx
movss %xmm0, (%rbp,%rcx,4)
incq %r12
cmpq %rax, %r12
jl 0x68e7
movl %eax, %edx
imull %r15d, %edx
movq 0x10(%rsp), %rsi
movl (%rsi,%r15,4), %ecx
movss (%rsp), %xmm1
jmp 0x6932
movl %eax, %edx
imull %r15d, %edx
addl %ecx, %edx
movslq %edx, %rcx
movss (%rbp,%rcx,4), %xmm0
addss 0x18a3(%rip), %xmm0     # 0x81e8
movss %xmm0, (%rbp,%rcx,4)
movl %eax, %r12d
incq %r15
movq (%r13), %rbx
movq 0x8(%r13), %rax
subq %rbx, %rax
sarq $0x2, %rax
movslq %r12d, %rcx
xorl %edx, %edx
divq %rcx
cmpq %r15, %rax
ja 0x67d5
xorps %xmm0, %xmm0
cvtsi2ssl 0x8(%rsp), %xmm0
divss %xmm0, %xmm1
movq 0x18(%r14), %rax
movss %xmm1, (%rax)
cmpb $0x0, 0xc(%rsp)
je 0x69be
movq 0x18(%r13), %rax
movq 0x20(%r13), %rcx
subq %rax, %rcx
je 0x69be
sarq $0x2, %rcx
cmpq $0x1, %rcx
adcq $0x0, %rcx
xorl %edx, %edx
movss (%rax,%rdx,4), %xmm1
divss %xmm0, %xmm1
movss %xmm1, (%rax,%rdx,4)
incq %rdx
cmpq %rdx, %rcx
jne 0x69a8
addq $0x28, %rsp
popq %rbx
popq %r12
popq %r13
popq %r14
popq %r15
popq %rbp
retq
nop