A quick summary of compiling and running GPT-2 (02/26/2022)

[wangkuiyi]

I used the string "Hello, my dog is cute" as a prompt and ran James' KV-caching-enabled GPT-2 implementation in the following ways. The expected result of the encoding should be [13], which in GPT-2 means "the period."

1. 🟢 Running the JAX program without jax.jit on an M1 Max MacBook Pro was smooth. The result is correct, [13].

2. 🟢 Compiling the JAX program into vmvx using IREE and running it using iree-run-module.

   python models/gpt2/export.py  # Creates /tmp/gpt2.mlir
   iree-compile /tmp/gpt2.mlir \
     --iree-hal-target-backends=vmvx \
     --iree-input-type=mhlo > /tmp/gpt2-vmvx.vmfb
   
   iree-run-module --module=/tmp/gpt2-vmvx.vmfb \
     --device="local-task" \
     --function=encode \
     --input="1x8xi32=[[15496 11 616 3290 318 13779 0 0]]" \
     --input="1xi32=[6]"

   The result is correct.

   EXEC @encode
   result[0]: hal.buffer_view
   1x1xi32=[13]
   

3. 🟡 Compiling the JAX program into metal code using this branch of IREE, running it using iree-run-module on an M1 Max MacBook Pro.

   iree-compile /tmp/gpt2.mlir \
     --iree-hal-target-backends=metal \
     --iree-input-type=mhlo  > /tmp/gpt2-metal.vmfb
   
   iree-run-module --module=/tmp/gpt2-metal.vmfb \
     --device=metal \
     --function=encode \
     --input="1x8xi32=[[15496 11 616 3290 318 13779 0 0]]" \
     --input="1xi32=[6]"

   It runs and retuns an answer. Unfortunately, the result is wrong (not 13):

   EXEC @encode
   result[0]: hal.buffer_view
   1x1xi32=[50271]
   

4. 🔴 Running the above /tmp/gpt2-metal.vmfb on an iPhone 13 with A chip. The iOS app crashed at allocating a tensor. Details are at here.

5. 🔴 Compiling the JAX program into llvm-cpu code on an M1 Max MacBook Pro and run it using iree-run-module on the same computer.

   iree-compile /tmp/gpt2.mlir \
     --iree-hal-target-backends=llvm-cpu \
     --iree-input-type=mhlo > /tmp/gpt2-cpu.vmfb
   iree-run-module --module=/tmp/gpt2-cpu.vmfb \
     --device="local-task" \
     --function=encode \
     --input="1x8xi32=[[15496 11 616 3290 318 13779 0 0]]" \
     --input="1xi32=[6]"

   The running of iree-run-module crashed.

   EXEC @encode
   Segmentation fault: 11
   

   I also tried to run this vmfb in the iOS Simulator. It also crashed. The team has been investigating this issue EXC_BAD_ACCESS signal received executing GPT2 llvm-cpu on iOS Simulator #12369

[vmurali]

Updates:

a) gpt2 works fine on x86_64 linux with or without the masking patch.

b) gpt2 works fine on arm64 linux using the current HEAD

c) gpt2 on Apple's M1+Metal - produces the wrong answer using this branch. Unfortunately, I cannot go back and forth on the commits because the patch for making IREE on Metal is still WIP by @antiagainst and the existing patches don't compose well.

d) gpt2 on Apple's M1 without Metal segfaults using the current HEAD. The standalone kernel responsible for segfaulting is this. Still trying to find the root cause of the segfault.

[vmurali]

Updates:

https://gist.github.com/afec50c1bab4b3a4c4cc786271d3ef7f contains a smaller kernel

The crashing instructions are as follows:

    0x10019c41c: add    x18, x11, x13, lsl #2 ; x18 set to 0x0000000101800100
    0x10019c420: mov    x0, #-0x4 ; x18 set to 0!!!
    0x10019c424: ldr    q0, [x18] ; segfault

Reading the ARM document, apparently x18 is a platform specific register.

X18 is the platform register and is reserved for the use of platform ABIs. This is an additional temporary register on platforms that don't assign a special meaning to it. from https://developer.arm.com/documentation/den0024/a/The-A64-instruction-set/System-control-and-other-instructions/Debug-instructions

https://gist.github.com/8eab36cd1eb39beb83b458eabbfc3913 is the assembly generated for the dispatch by iree

This may not be the only error though (I don't see this error for 768x13 - it shows up only for 8x13. Consistently though.

[vmurali]

Changing
https://cs.github.com/llvm/llvm-project/blob/e65e7b27fdcc286c799e748fa51dccdf39321ee3/llvm/lib/TargetParser/AArch64TargetParser.cpp#L98 function to return true always fixes the bug (the whole gpt2 model works now). Somehow TT.isOSDarwinn() returns false right now.

[vmurali]

Fixed for llvm-cpu backend with #12460