[spirv] Incorrect int4 vicuna model output

[yzhang93]

What happened?

Vicuna int4 model outputs are different compared to fake_int8 model (int4 weights in int8 container). The results begin diverging from dispatch 5.

Build is based on afc8705

Steps to reproduce your issue

1. Download model second_vicuna_int4.mlir.
2. Commands to reproduce:
   ./iree-compile --iree-input-type=none --iree-vm-bytecode-module-output-format=flatbuffer-binary --iree-hal-target-backends=vulkan --mlir-print-debuginfo --mlir-print-op-on-diagnostic=false --iree-llvmcpu-target-cpu-features=host --iree-stream-resource-index-bits=64 --iree-vm-target-index-bits=64 --iree-vm-bytecode-module-strip-source-map=false --iree-util-zero-fill-elided-attrs --iree-vm-target-truncate-unsupported-floats --iree-codegen-check-ir-before-llvm-conversion=false --iree-vulkan-target-triple=rdna3-unknown-linux --iree-opt-const-expr-hoisting=false --iree-consteval-jit-debug=true ~/Downloads/second_vicuna_int4.mlir -o vicuna_vulkan_i4.vmfb --iree-stream-resource-max-allocation-size=3221225472 --iree-flow-break-dispatch=@forward:5

./iree-run-module --device_allocator=caching --vulkan_vma_allocator=false --module=vicuna_vulkan_i4.vmfb --device=vulkan --function=forward --input=1x1xi64 --input=1x32x1x128xf32 --input=1x32x1x128xf32 --input=1x32x1x128xf32 --input=1x32x1x128xf32 --input=1x32x1x128xf32 --input=1x32x1x128xf32 --input=1x32x1x128xf32 --input=1x32x1x128xf32 --input=1x32x1x128xf32 --input=1x32x1x128xf32 --input=1x32x1x128xf32 --input=1x32x1x128xf32 --input=1x32x1x128xf32 --input=1x32x1x128xf32 --input=1x32x1x128xf32 --input=1x32x1x128xf32 --input=1x32x1x128xf32 --input=1x32x1x128xf32 --input=1x32x1x128xf32 --input=1x32x1x128xf32 --input=1x32x1x128xf32 --input=1x32x1x128xf32 --input=1x32x1x128xf32 --input=1x32x1x128xf32 --input=1x32x1x128xf32 --input=1x32x1x128xf32 --input=1x32x1x128xf32 --input=1x32x1x128xf32 --input=1x32x1x128xf32 --input=1x32x1x128xf32 --input=1x32x1x128xf32 --input=1x32x1x128xf32 --input=1x32x1x128xf32 --input=1x32x1x128xf32 --input=1x32x1x128xf32 --input=1x32x1x128xf32 --input=1x32x1x128xf32 --input=1x32x1x128xf32 --input=1x32x1x128xf32 --input=1x32x1x128xf32 --input=1x32x1x128xf32 --input=1x32x1x128xf32 --input=1x32x1x128xf32 --input=1x32x1x128xf32 --input=1x32x1x128xf32 --input=1x32x1x128xf32 --input=1x32x1x128xf32 --input=1x32x1x128xf32 --input=1x32x1x128xf32  --input=1x32x1x128xf32 --input=1x32x1x128xf32 --input=1x32x1x128xf32 --input=1x32x1x128xf32 --input=1x32x1x128xf32 --input=1x32x1x128xf32 --input=1x32x1x128xf32 --input=1x32x1x128xf32 --input=1x32x1x128xf32 --input=1x32x1x128xf32 --input=1x32x1x128xf32 --input=1x32x1x128xf32 --input=1x32x1x128xf32 --input=1x32x1x128xf32 --input=1x32x1x128xf32 --output=@vulkan_dispatch5.npy

3. The results I got for dispatch 5 from the above commands: https://storage.googleapis.com/shark-public/vivian/vicuna_i4/vulkan_dispatch5.npy
4. Golden results for comparison: https://storage.googleapis.com/shark-public/vivian/vicuna_i4/vulkan_dispatch5_golden.npy

[yzhang93]

With emulation pass added #14737, I got compilation error
error: failed to legalize operation 'vector.transfer_read' that was explicitly marked illegal

Here's dump IR after all: https://storage.googleapis.com/shark-public/vivian/vicuna_i4/vicuna_vulkan_dump_after_all.txt

[MaheshRavishankar]

With emulation pass added #14737, I got compilation error error: failed to legalize operation 'vector.transfer_read' that was explicitly marked illegal

Here's dump IR after all: https://storage.googleapis.com/shark-public/vivian/vicuna_i4/vicuna_vulkan_dump_after_all.txt

I think you need to pattern to handle vector.transfer_read (similar to vector.load).

[yzhang93]

@MaheshRavishankar @antiagainst My concern is now the problem seems to have changed from what we saw on CPU, since the results are the same w/wo -optimize-bindings flag, but still not correct compared to fake int8. Not sure if the emulation pass (with vector.transfer_read pattern added) we used on the CPU is the route to fix this problem.

[MaheshRavishankar]

My suspicion is that the issues are similar but are buried in the spirv conversion. Doing this ahead of time head those out earlier and handle them correctly.

[yzhang93]

From the dump IR, I noticed it's trying to emulate i4 to i128 from the below logic? Are we still emulating i4 to i8? What should be the correct logic for vector.transfer_read emulation then?

%6 = affine.apply affine_map<()[s0] -> (s0 floordiv 32)>()[%0] loc(callsite("/home/vivian/Downloads/vicuna_unsharded_mlir_second_vicuna_int4.mlir":915:11 at "/home/vivian/Downloads/vicuna_unsharded_mlir_second_vicuna_int4.mlir":26:3))
%7 = affine.apply affine_map<()[s0] -> (s0 * 4 - (s0 floordiv 32) * 128)>()[%0] loc(callsite("/home/vivian/Downloads/vicuna_unsharded_mlir_second_vicuna_int4.mlir":915:11 at "/home/vivian/Downloads/vicuna_unsharded_mlir_second_vicuna_int4.mlir":26:3))
%8 = vector.transfer_read %1[%workgroup_id_x, %6, %7], %c0_i4 {in_bounds = [true]} : memref<4096x32x128xi4, strided<[4096, 128, 1], offset: 1050705920>, #hal.descriptor_type<storage_buffer>>, vector<4xi4> loc(callsite("/home/vivian/Downloads/vicuna_unsharded_mlir_second_vicuna_int4.mlir":915:11 at "/home/vivian/Downloads/vicuna_unsharded_mlir_second_vicuna_int4.mlir":26:3))

[MaheshRavishankar]

From the dump IR, I noticed it's trying to emulate i4 to i128 from the below logic? Are we still emulating i4 to i8? What should be the correct logic for vector.transfer_read emulation then?

%6 = affine.apply affine_map<()[s0] -> (s0 floordiv 32)>()[%0] loc(callsite("/home/vivian/Downloads/vicuna_unsharded_mlir_second_vicuna_int4.mlir":915:11 at "/home/vivian/Downloads/vicuna_unsharded_mlir_second_vicuna_int4.mlir":26:3))
%7 = affine.apply affine_map<()[s0] -> (s0 * 4 - (s0 floordiv 32) * 128)>()[%0] loc(callsite("/home/vivian/Downloads/vicuna_unsharded_mlir_second_vicuna_int4.mlir":915:11 at "/home/vivian/Downloads/vicuna_unsharded_mlir_second_vicuna_int4.mlir":26:3))
%8 = vector.transfer_read %1[%workgroup_id_x, %6, %7], %c0_i4 {in_bounds = [true]} : memref<4096x32x128xi4, strided<[4096, 128, 1], offset: 1050705920>, #hal.descriptor_type<storage_buffer>>, vector<4xi4> loc(callsite("/home/vivian/Downloads/vicuna_unsharded_mlir_second_vicuna_int4.mlir":915:11 at "/home/vivian/Downloads/vicuna_unsharded_mlir_second_vicuna_int4.mlir":26:3))

My thought is we emulate to i8 for now and the spirv lowering handles the lowering of i8. Hope that fixes the issue.

Not sure why it is emulating to i28

[antiagainst]

So diff'ing the data, it seems that we have an element ordering issue:

>>> import numpy as np
>>> data1 = np.load('vulkan_dispatch5.npy')
>>> data0 = np.load('vulkan_dispatch5_golden.npy')
>>> np.set_printoptions(linewidth=100)
>>> print(data0[0,:32])
[-0.01193237 -0.02619425  0.00946045 -0.01906331 -0.02619425  0.00946045  0.00946045  0.00232951
 -0.00480143 -0.00480143 -0.00480143  0.00946045 -0.00480143 -0.01193237  0.00946045 -0.01193237
 -0.00480143 -0.00480143 -0.04045613  0.03085327 -0.00480143 -0.01906331  0.01659139  0.01659139
 -0.00480143  0.01659139 -0.00480143 -0.00480143 -0.00480143  0.00232951 -0.01906331  0.02372233]
>>> print(data1[0,:32])
[ 0.00232951  0.00946045  0.00946045 -0.02619425 -0.01906331  0.00946045 -0.02619425 -0.01193237
 -0.01193237  0.00946045 -0.01193237 -0.00480143  0.00946045 -0.00480143 -0.00480143 -0.00480143
  0.01659139  0.01659139 -0.01906331 -0.00480143  0.03085327 -0.04045613 -0.00480143 -0.00480143
  0.02372233 -0.01906331  0.00232951 -0.00480143 -0.00480143 -0.00480143  0.01659139 -0.00480143]

For each of the 8 elements, the vulkan one is the reverse of the golden file. We have some inconsistency regarding element ordering/endianness and bitcasting/etc..

[stellaraccident]

Looks like wrong order of nibbles in the i32? That's actually somewhat reassuring if that is the issue.

[yzhang93]

For each of the 8 elements, the vulkan one is the reverse of the golden file. We have some inconsistency regarding element ordering/endianness and bitcasting/etc..

Nice finding, Thanks @antiagainst!

[antiagainst]

Right, it's my oversight when throwing in a new pattern previously. :( Uploaded #14755 to fix this particular dispatch. Don't know whether this fixes all the issues though.. Needs verification.