[all-commits] [llvm/llvm-project] 28fef9: [mlir][VectorOps] Fix folding of vector.extract fr...

[Benjamin Maxwell via All-commits]

  Branch: refs/heads/main
  Home:   https://github.com/llvm/llvm-project
  Commit: 28fef902fc9f872fcdcd95237f425f58c9609ab4
      https://github.com/llvm/llvm-project/commit/28fef902fc9f872fcdcd95237f425f58c9609ab4
  Author: Benjamin Maxwell <benjamin.maxwell at arm.com>
  Date:   2023-08-04 (Fri, 04 Aug 2023)

  Changed paths:
    M mlir/lib/Dialect/Vector/IR/VectorOps.cpp
    M mlir/test/Dialect/Vector/canonicalize.mlir

  Log Message:
  -----------
  [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