This is an archive of the discontinued LLVM Phabricator instance.

[mlir][VectorOps] Fix folding of vector.extract from stretch vector.broadcast
ClosedPublic

Authored by benmxwl-arm on Aug 3 2023, 7:26 AM.

Details

Summary

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.

Diff Detail

Event Timeline

benmxwl-arm created this revision.Aug 3 2023, 7:26 AM
Herald added a project: Restricted Project. · View Herald Transcript
benmxwl-arm requested review of this revision.Aug 3 2023, 7:26 AM
benmxwl-arm edited the summary of this revision. (Show Details)Aug 3 2023, 7:44 AM
dcaballe accepted this revision.Aug 3 2023, 9:29 AM

Thanks!

This revision is now accepted and ready to land.Aug 3 2023, 9:29 AM
awarzynski accepted this revision.Aug 3 2023, 1:36 PM

LGTM, thanks!