[mlir][VectorOps] Fix folding of vector.extract from stretch vector.broadcast
Previously, foldExtractFromBroadcast() would incorrectly fold: func.func @extract_from_stretch_broadcast(%src: vector<3x1x2xf32>) -> f32 { %0 = vector.broadcast %src : vector<3x1x2xf32> to vector<3x4x2xf32> %1 = vector.extract %0[0, 2, 0] : vector<3x4x2xf32> return %1: f32 } to: func.func @extract_from_stretch_broadcast(%src: vector<3x1x2xf32>) -> f32 { %0 = vector.extract %src[0, 2, 0] : vector<3x1x2xf32> return %0: f32 } This was due to the wrong offset being used when zeroing the "dim-1" broadcasted dims. It should use the difference in rank across the broadcast as the starting offset, as the ranks after that are the ones that could have been stretched. Reviewed By: awarzynski, dcaballe Differential Revision: https://reviews.llvm.org/D157003
Loading
Please sign in to comment