search

ROCm / AMDMIGraphX

AMD's graph optimization engine.
https://rocm.docs.amd.com/projects/AMDMIGraphX/en/latest/
MIT License
185 stars 86 forks source link

Prefill fix for multiple output cases #3605

Closed shivadbhavsar closed 4 days ago

shivadbhavsar commented 1 week ago

Background: If we request a mlir kernel with multiple outputs, the output buffer will be allocated as a tuple. Example IR before compile_ops pass:

module: "main"
output = @param:output -> float_type, {2, 4, 64, 64}, {16384, 4096, 64, 1}
bias = @param:bias -> float_type, {40960}, {1}
scale = @param:scale -> float_type, {40960}, {1}
@3 = @literal{ ... } -> float_type, {320}, {1}
w2 = @param:w2 -> float_type, {4, 320, 3, 3}, {2880, 9, 3, 1}
w1 = @param:w1 -> float_type, {320, 4, 3, 3}, {36, 9, 3, 1}
x = @param:x -> float_type, {2, 4, 64, 64}, {16384, 4096, 64, 1}
@7 = reshape_lazy[dims={32, 10}](@3) -> float_type, {32, 10}, {10, 1}
@8 = broadcast[axis=1,out_lens={2, 32, 10, 64, 64}](@7) -> float_type, {2, 32, 10, 64, 64}, {0, 10, 1, 0, 0}
@9 = hip::allocate[shape=[float_type, {2, 32, 1, 1, 1}, {32, 1, 1, 1, 1}, float_type, {2, 32, 1, 1, 1}, {32, 1, 1, 1, 1}, float_type, {2, 32, 10, 64, 64}, {1310720, 40960, 4096, 64, 1}]] -> [float_type, {2, 32, 1, 1, 1}, {32, 1, 1, 1, 1}, float_type, {2, 32, 1, 1, 1}, {32, 1, 1, 1, 1}, float_type, {2, 32, 10, 64, 64}, {1310720, 40960, 4096, 64, 1}]
@10 = gpu::precompile_op[op=gpu::mlir_op[op=gpu::mlir_op[op=convolution[padding={1, 1, 1, 1},stride={1, 1},dilation={1, 1},group=1,padding_mode=0]]],additional_args=1,ignore_modules=0,output_shape=nullopt](@8,x,w1,@9), [mlir_main:pointwise0_main:pointwise0:main:pointwise3:main:reduce_sum1:main:pointwise5:main:pointwise2:main:pointwise1:main:reduce_sum0_reshape_split] -> [float_type, {2, 32, 1, 1, 1}, {32, 1, 1, 1, 1}, float_type, {2, 32, 1, 1, 1}, {32, 1, 1, 1, 1}, float_type, {2, 32, 10, 64, 64}, {1310720, 40960, 4096, 64, 1}]
@11 = get_tuple_elem[index=2](@10) -> float_type, {2, 32, 10, 64, 64}, {1310720, 40960, 4096, 64, 1}
@12 = get_tuple_elem[index=0](@10) -> float_type, {2, 32, 1, 1, 1}, {32, 1, 1, 1, 1}
@13 = get_tuple_elem[index=1](@10) -> float_type, {2, 32, 1, 1, 1}, {32, 1, 1, 1, 1}
@14 = gpu::precompile_op[op=gpu::mlir_op[op=convolution[padding={1, 1, 1, 1},stride={1, 1},dilation={1, 1},group=1,padding_mode=0]],additional_args=1,ignore_modules=0,output_shape=nullopt](@12,@13,@11,scale,bias,w2,output), [main:pointwise5:mlir_convolution1] -> float_type, {2, 4, 64, 64}, {16384, 4096, 64, 1}

Bug: Our code is not able to prefill a buffer in the multi output case

Fix: MLIR doesnt accept tuples and so we flatten them when giving inputs to mlir code objects. Same type of flattening needs to happen when determining which buffers in the allocated tuple need to be prefilled.

Assumptions:

If either of those assumptions become untrue in the future, we will need better logic to hand this.

shivadbhavsar commented 1 week ago

I included the MLIR bump from #3586 since thats where we first found this bug.

codecov[bot] commented 1 week ago

Codecov Report

All modified and coverable lines are covered by tests :white_check_mark:

Project coverage is 92.17%. Comparing base (5388984) to head (d7a122f). Report is 9 commits behind head on develop.

Additional details and impacted files ```diff @@ Coverage Diff @@ ## develop #3605 +/- ## ========================================= Coverage 92.16% 92.17% ========================================= Files 512 513 +1 Lines 21412 21536 +124 ========================================= + Hits 19735 19851 +116 - Misses 1677 1685 +8 ```

:umbrella: View full report in Codecov by Sentry.
:loudspeaker: Have feedback on the report? Share it here.

migraphx-bot commented 5 days ago
Test Batch Rate new
d7a122		 Rate old
f5df00		 Diff Compare
torchvision-resnet50 64 3,257.99 3,256.17 0.06% :white_check_mark:
torchvision-resnet50_fp16 64 6,984.89 6,985.11 -0.00% :white_check_mark:
torchvision-densenet121 32 2,435.53 2,435.96 -0.02% :white_check_mark:
torchvision-densenet121_fp16 32 4,078.53 4,058.51 0.49% :white_check_mark:
torchvision-inceptionv3 32 1,638.00 1,636.64 0.08% :white_check_mark:
torchvision-inceptionv3_fp16 32 2,763.12 2,762.67 0.02% :white_check_mark:
cadene-inceptionv4 16 776.31 776.18 0.02% :white_check_mark:
cadene-resnext64x4 16 811.16 811.76 -0.07% :white_check_mark:
slim-mobilenet 64 7,534.78 7,534.07 0.01% :white_check_mark:
slim-nasnetalarge 64 211.44 211.45 -0.01% :white_check_mark:
slim-resnet50v2 64 3,505.37 3,504.24 0.03% :white_check_mark:
bert-mrpc-onnx 8 1,152.93 1,149.67 0.28% :white_check_mark:
bert-mrpc-tf 1 473.06 463.86 1.98% :white_check_mark:
pytorch-examples-wlang-gru 1 417.22 420.46 -0.77% :white_check_mark:
pytorch-examples-wlang-lstm 1 476.69 381.56 24.93% :high_brightness:
torchvision-resnet50_1 1 768.24 780.85 -1.61% :white_check_mark:
cadene-dpn92_1 1 412.33 405.55 1.67% :white_check_mark:
cadene-resnext101_1 1 382.45 383.55 -0.29% :white_check_mark:
onnx-taau-downsample 1 342.55 343.07 -0.15% :white_check_mark:
dlrm-criteoterabyte 1 33.31 33.34 -0.09% :white_check_mark:
dlrm-criteoterabyte_fp16 1 52.70 52.74 -0.08% :white_check_mark:
agentmodel 1 8,258.46 8,306.27 -0.58% :white_check_mark:
unet_fp16 2 58.83 58.82 0.00% :white_check_mark:
resnet50v1_fp16 1 962.75 1,001.66 -3.88% :red_circle:
resnet50v1_int8 1 1,011.54 995.76 1.59% :white_check_mark:
bert_base_cased_fp16 64 1,169.90 1,171.04 -0.10% :white_check_mark:
bert_large_uncased_fp16 32 363.52 363.62 -0.03% :white_check_mark:
bert_large_fp16 1 200.45 198.87 0.79% :white_check_mark:
distilgpt2_fp16 16 2,203.83 2,204.83 -0.05% :white_check_mark:
yolov5s 1 533.13 540.84 -1.42% :white_check_mark:
tinyllama 1 43.70 43.47 0.53% :white_check_mark:
vicuna-fastchat 1 176.20 176.64 -0.25% :white_check_mark:
whisper-tiny-encoder 1 418.15 418.46 -0.07% :white_check_mark:
whisper-tiny-decoder 1 425.71 433.85 -1.88% :white_check_mark:

This build is not recommended to merge :red_circle:

migraphx-bot commented 5 days ago


     :white_check_mark: bert-mrpc-onnx: PASSED: MIGraphX meets tolerance
     :white_check_mark: bert-mrpc-tf: PASSED: MIGraphX meets tolerance
     :white_check_mark: pytorch-examples-wlang-gru: PASSED: MIGraphX meets tolerance
     :white_check_mark: pytorch-examples-wlang-lstm: PASSED: MIGraphX meets tolerance
     :white_check_mark: torchvision-resnet50_1: PASSED: MIGraphX meets tolerance
     :white_check_mark: cadene-dpn92_1: PASSED: MIGraphX meets tolerance
     :white_check_mark: cadene-resnext101_1: PASSED: MIGraphX meets tolerance
     :white_check_mark: dlrm-criteoterabyte: PASSED: MIGraphX meets tolerance
     :white_check_mark: agentmodel: PASSED: MIGraphX meets tolerance
     :white_check_mark: unet: PASSED: MIGraphX meets tolerance
     :white_check_mark: resnet50v1: PASSED: MIGraphX meets tolerance
     :white_check_mark: bert_base_cased_fp16: PASSED: MIGraphX meets tolerance
:red_circle:bert_large_uncased_fp16: FAILED: MIGraphX is not within tolerance - check verbose output

     :white_check_mark: bert_large: PASSED: MIGraphX meets tolerance
     :white_check_mark: yolov5s: PASSED: MIGraphX meets tolerance
     :white_check_mark: tinyllama: PASSED: MIGraphX meets tolerance
     :white_check_mark: vicuna-fastchat: PASSED: MIGraphX meets tolerance
     :white_check_mark: whisper-tiny-encoder: PASSED: MIGraphX meets tolerance
     :white_check_mark: whisper-tiny-decoder: PASSED: MIGraphX meets tolerance
     :white_check_mark: distilgpt2_fp16: PASSED: MIGraphX meets tolerance