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

[benmxwl-arm]

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.

[dcaballe]

Thanks!

[awarzynski]

LGTM, thanks!

mlir/lib/Dialect/Vector/IR/VectorOps.cpp

static Value foldExtractFromBroadcast(ExtractOp extractOp) {
   auto extractVecType = llvm::dyn_cast<VectorType>(extractOp.getType());
   auto broadcastVecType = llvm::dyn_cast<VectorType>(source.getType());
   if (extractVecType && broadcastVecType &&
       extractVecType.getShape() !=
           broadcastVecType.getShape().take_back(extractResultRank))
     return Value();
 
   auto broadcastOp = cast<vector::BroadcastOp>(defOp);
-  int64_t rankDiff = broadcastSrcRank - extractResultRank;
+  int64_t broadcastDstRank = broadcastOp.getResultVectorType().getRank();
+
   // Detect all the positions that come from "dim-1" broadcasting.
   // These dimensions correspond to "dim-1" broadcasted dims; set the mathching
   // extract position to `0` when extracting from the source operand.
   llvm::SetVector<int64_t> broadcastedUnitDims =
       broadcastOp.computeBroadcastedUnitDims();
   SmallVector<int64_t> extractPos(extractOp.getPosition());
-  for (int64_t i = rankDiff, e = extractPos.size(); i < e; ++i)
+  int64_t broadcastRankDiff = broadcastDstRank - broadcastSrcRank;
+  for (int64_t i = broadcastRankDiff, e = extractPos.size(); i < e; ++i)
     if (broadcastedUnitDims.contains(i))
       extractPos[i] = 0;
   // `rankDiff` leading dimensions correspond to new broadcasted dims, drop the
   // matching extract position when extracting from the source operand.
+  int64_t rankDiff = broadcastSrcRank - extractResultRank;
   extractPos.erase(extractPos.begin(),
                    std::next(extractPos.begin(), extractPos.size() - rankDiff));
   // OpBuilder is only used as a helper to build an I64ArrayAttr.
   OpBuilder b(extractOp.getContext());
   extractOp.setOperand(source);
   extractOp.setPosition(extractPos);
   return extractOp.getResult();
 }
void vector::TransposeOp::build(OpBuilder &builder, OperationState &result,
   result.addOperands(vector);
   result.addTypes(VectorType::get(transposedShape, vt.getElementType()));
   result.addAttribute(TransposeOp::getTranspAttrName(result.name),
                       builder.getI64ArrayAttr(transp));
 }
 
 OpFoldResult vector::TransposeOp::fold(FoldAdaptor adaptor) {
   // Eliminate splat constant transpose ops.
-  if (auto attr = llvm::dyn_cast_if_present<DenseElementsAttr>(adaptor.getVector()))
+  if (auto attr =
+          llvm::dyn_cast_if_present<DenseElementsAttr>(adaptor.getVector()))
     if (attr.isSplat())
       return attr.reshape(getResultVectorType());
 
   // Eliminate identity transpose ops. This happens when the dimensions of the
   // input vector remain in their original order after the transpose operation.
   SmallVector<int64_t, 4> transp;
   getTransp(transp);
 

mlir/test/Dialect/Vector/canonicalize.mlir

func.func @extract_from_broadcast(%src: vector<1x1x1xf32>) -> vector<1xf32> {
   %0 = vector.broadcast %src : vector<1x1x1xf32> to vector<1x1x32x1xf32>
 
   //  CHECK-NEXT:   %0 = vector.extract {{.*}}[0, 0] : vector<1x1x1xf32>
   //  CHECK-NEXT:   return %0 : vector<1xf32>
   %1 = vector.extract %0[0, 0, 31] : vector<1x1x32x1xf32>
   return %1: vector<1xf32>
 }
 
+// CHECK-LABEL: func.func @extract_from_stretch_broadcast
+func.func @extract_from_stretch_broadcast(%src: vector<3x1x2xf32>) -> f32 {
+  //  CHECK-NEXT:  %0 = vector.extract {{.*}}[0, 0, 0] : vector<3x1x2xf32>
+  //  CHECK-NEXT:  return %0 : f32
+  %0 = vector.broadcast %src : vector<3x1x2xf32> to vector<3x4x2xf32>
+  %1 = vector.extract %0[0, 2, 0] : vector<3x4x2xf32>
+  return %1: f32
+}
+
 // -----
 // CHECK-LABEL: func.func @extract_strided_slice_of_constant_mask
 func.func @extract_strided_slice_of_constant_mask() -> vector<5x7xi1>{
   //  CHECK-NEXT:   %[[RES:.*]] = vector.constant_mask [5, 4] : vector<5x7xi1>
   //  CHECK-NEXT:   return %[[RES]] : vector<5x7xi1>
   %c4 = arith.constant 4 : index
   %c10 = arith.constant 10 : index
   %mask = vector.create_mask %c10, %c4 : vector<12x7xi1>