[Codegen] Fix ArgCompare vectorization

[bangtianliu]

This PR fixes ArgCompare vectorization to use input shape instead of the output shape for inferring vector sizes.

The existing implementation uses inferSizesFromIR(argCompareOp.getDpsInits()[0]), which fails for rank-reducing operations. Inferring from the output returns incorrect sizes, especially when outputs come from rank-reducing tensor.extract_slice operations.

Example failure case:

 %extracted_slice_2 = tensor.extract_slice %extracted_slice[0] [1] [1] : tensor<1xf32> to tensor<f32>
 %32:2 = iree_linalg_ext.arg_compare dimension(0)
     ins(%28, %31 : tensor<1024xf32>, tensor<1024xi64>)
     outs(%extracted_slice_2, %extracted_slice_4 : tensor<f32>, tensor<i64>)

For reduction ops, inferSizesFromIR(linalgOp, ...) uses linalgOp.getNumLoops() to map ALL loop dimensions (parallel + reduction) to operands.

Assisted-by: Claude Code