x86_64 nasm jumps to the wrong location

I'm working on code to enumerate the PCI bus, but have found that the jz statement for the loop over each device jumps to the wrong location (not even a label). The register function should be getting called for each time cmp ax, 0xffff is inequal, which should be more than once. It is only getting called once.

register:                               ; eax = edi = config offset of the function
    mov dx, ADDR_PRT
    add eax, 0x08
    in eax, dx
    shr eax, 16
    mov dx, ax
    call checkpoint
    mov eax, edi
rmsd:   cmp dx, 0x0601                  ; mass storage devices
    je ahci_register                ; register an AHCI controller
    ret                             ; couldn't find it, ignore it


pci_init:
    mov edi, 0x80000000
ilp0:   mov rax, rdi
    mov dx, ADDR_PRT
    out dx, eax
    mov dx, DATA_PRT
    in eax, dx
    cmp ax, 0xffff
    je ilp0c0
    push rdi
    mov rax, rdi
    call register
    pop rdi
ilp0c0: add rdi, 0x100
    test edi, 0xff000000            ; code jupms to the line before this
    jz ilp0
    ret

Code is assembled as a PE file and then linked using lld-link and run using EFI.