Prefill fix for multiple output cases

[shivadbhavsar]

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:

• Inputs will never be tuples
• Allocated buffers are always the last input (which i believe is generally true everywhere in migraphx?)

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

[shivadbhavsar]

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

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

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

Additional details and impacted files

@@            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     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[migraphx-bot]

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

This build is not recommended to merge 🔴

[migraphx-bot]

     ✅ bert-mrpc-onnx: PASSED: MIGraphX meets tolerance

     ✅ bert-mrpc-tf: PASSED: MIGraphX meets tolerance

     ✅ pytorch-examples-wlang-gru: PASSED: MIGraphX meets tolerance

     ✅ pytorch-examples-wlang-lstm: PASSED: MIGraphX meets tolerance

     ✅ torchvision-resnet50_1: PASSED: MIGraphX meets tolerance

     ✅ cadene-dpn92_1: PASSED: MIGraphX meets tolerance

     ✅ cadene-resnext101_1: PASSED: MIGraphX meets tolerance

     ✅ dlrm-criteoterabyte: PASSED: MIGraphX meets tolerance

     ✅ agentmodel: PASSED: MIGraphX meets tolerance

     ✅ unet: PASSED: MIGraphX meets tolerance

     ✅ resnet50v1: PASSED: MIGraphX meets tolerance

     ✅ bert_base_cased_fp16: PASSED: MIGraphX meets tolerance

🔴bert_large_uncased_fp16: FAILED: MIGraphX is not within tolerance - check verbose output

     ✅ bert_large: PASSED: MIGraphX meets tolerance

     ✅ yolov5s: PASSED: MIGraphX meets tolerance

     ✅ tinyllama: PASSED: MIGraphX meets tolerance

     ✅ vicuna-fastchat: PASSED: MIGraphX meets tolerance

     ✅ whisper-tiny-encoder: PASSED: MIGraphX meets tolerance

     ✅ whisper-tiny-decoder: PASSED: MIGraphX meets tolerance

     ✅ distilgpt2_fp16: PASSED: MIGraphX meets tolerance