Skip to content
Commit 28fef902 authored by Benjamin Maxwell's avatar Benjamin Maxwell
Browse files

[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
parent fcb0294b
Loading
Loading
Loading
Loading
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please to comment