[mlir][vector] Fold scalar vector.extract of non-splat n-D constants

[kuhar]

Summary

Add a new pattern to fold vector.extract over n-D constants that extract scalars.
The previous code handled ND splat constants only. The new pattern is conservative and does handle sub-vector constants.

This is to aid the arith::EmulateWideInt pass which emits a lot of 2-element vector constants.

[dcaballe]

Thanks, @kuhar! LGTM. Just a few comments!

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

1575–1576

I think this condition is guaranteed by definition of the vector type.

1577–1578

Would it be too much work to also add support for n-D cases where a scalar is extracted? It would be very useful in general! Having it working for n-D would be great if it's just about adding some extra logic to compute the position. Supporting a sub-vector extraction would require more work that we can add in the future.

1587

I wonder if it would make sense to merge this pattern with the one above. They current share some code and when we eventually support n-D cases, the would be even more similar, except by the way the constant value is extracted, right?

[kuhar]

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

1575–1576

Thanks, I wanted to exclude scalable vectors, but now I see I used the wrong predicate for this.

1577–1578

I considered this but was warned by @antiagainst to be careful not to blow up the size of the IR. Thinking about it more, I think this would be fine if the extracted element was a scalar instead of a nested vector constant, because we would always replace one vector.extract with one arith.constant with a single APInt/APFloat.

IE I think it would be safe to support n-D cases when all positions are provided, but not when some dims are still left. Does this make sense?

[kuhar]

Fixed scalable vector check

[Mogball]

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

1570

Please match this using m_Constant. The defining op might be a different kind of constant op.

[dcaballe]

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

1577–1578

IE I think it would be safe to support n-D cases when all positions are provided, but not when some dims are still left. Does this make sense?

Exactly! That's what I meant.

Regarding the sub-vector extraction, I'm not sure I understand the IR size concern. Could you please elaborate? This folding is something that will eventually happen in LLVM so I think the sooner we expose these constants the better.

[kuhar]

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

1577–1578

Regarding the sub-vector extraction, I'm not sure I understand the IR size concern. Could you please elaborate? This folding is something that will eventually happen in LLVM so I think the sooner we expose these constants the better.

Say you have a large n-D constant vector: %v = arith.constant dence<...> : vector<...> and perform a series of extractions for sub-vectors.
If you encode it with vector.extract, each extraction needs < n * 64-bit positions. If you convert it to a series of constants, each representing some sub-vector, these constants may be much larger than the positions. I don't know if the in-memory attributes will be deduplicated or not, but this seems like something that can negatively affect the printing/parsing time in pathological cases.

[kuhar]

Handle n-D vectors. Match constants vectors.

[kuhar]

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

1587

IMO the new pattern is sufficiently different that merging the two does not make the code cleaner.

[dcaballe]

Thanks! LGTM

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

1577–1578

If you convert it to a series of constants, each representing some sub-vector, these constants may be much larger than the positions. I don't know if the in-memory attributes will be deduplicated or not, but this seems like something that can negatively affect the printing/parsing time in pathological cases.

It's ok for me if we don't want to support the sub-vector cases for now but I don't think making printing/parsing slower would justify not introducing that canonicalization. Actually, splitting a massive constant into smaller ones could even help. I would expect identical constants to be deduplicated. Again, I think this will happen in LLVM anyways.

[Mogball]

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

1584

Can you drop a TODO to support other element types, like complex values?

[kuhar]

Added a TODO to support more types. Rebased.

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

if (extractResultRank < broadcastSrcRank)
       return failure();
     rewriter.replaceOpWithNewOp<vector::BroadcastOp>(
         extractOp, extractOp.getType(), source);
     return success();
   }
 };
 
 // Pattern to rewrite a ExtractOp(splat ConstantOp) -> ConstantOp.
-class ExtractOpConstantFolder final : public OpRewritePattern<ExtractOp> {
+class ExtractOpSplatConstantFolder final : public OpRewritePattern<ExtractOp> {
 public:
   using OpRewritePattern::OpRewritePattern;
 
   LogicalResult matchAndRewrite(ExtractOp extractOp,
                                 PatternRewriter &rewriter) const override {
-    // Return if 'extractStridedSliceOp' operand is not defined by a
+    // Return if 'ExtractOp' operand is not defined by a splat vector
     // ConstantOp.
-    auto constantOp = extractOp.getVector().getDefiningOp<arith::ConstantOp>();
-    if (!constantOp)
+    Value sourceVector = extractOp.getVector();
+    Attribute vectorCst;
+    if (!matchPattern(sourceVector, m_Constant(&vectorCst)))
       return failure();
-    auto dense = constantOp.getValue().dyn_cast<SplatElementsAttr>();
-    if (!dense)
+    auto splat = vectorCst.dyn_cast<SplatElementsAttr>();
+    if (!splat)
       return failure();
-    Attribute newAttr = dense.getSplatValue<Attribute>();
+    Attribute newAttr = splat.getSplatValue<Attribute>();
     if (auto vecDstType = extractOp.getType().dyn_cast<VectorType>())
       newAttr = DenseElementsAttr::get(vecDstType, newAttr);
     rewriter.replaceOpWithNewOp<arith::ConstantOp>(extractOp, newAttr);
     return success();
   }
 };
 
+// Pattern to rewrite a ExtractOp(vector<...xT> ConstantOp)[...] -> ConstantOp,
+// where the position array specifies a scalar element.
+class ExtractOpScalarVectorConstantFolder final
+    : public OpRewritePattern<ExtractOp> {
+public:
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(ExtractOp extractOp,
+                                PatternRewriter &rewriter) const override {
+    // Return if 'ExtractOp' operand is not defined by a compatible vector
+    // ConstantOp.
+    Value sourceVector = extractOp.getVector();
+    Attribute vectorCst;
+    if (!matchPattern(sourceVector, m_Constant(&vectorCst)))
+      return failure();
+
+    auto vecTy = sourceVector.getType().cast<VectorType>();
+    Type elemTy = vecTy.getElementType();
+    ArrayAttr positions = extractOp.getPosition();
+    if (vecTy.isScalable())
+      return failure();
+    // Do not allow extracting sub-vectors to limit the size of the generated
+    // constants.
+    if (vecTy.getRank() != static_cast<int64_t>(positions.size()))
+      return failure();
+    // TODO: Handle more element types, e.g., complex values.
+    if (!elemTy.isIntOrIndexOrFloat())
+      return failure();
+
+    // The splat case is handled by `ExtractOpSplatConstantFolder`.
+    auto dense = vectorCst.dyn_cast<DenseElementsAttr>();
+    if (!dense || dense.isSplat())
+      return failure();
+
+    // Calculate the flattened position.
+    int64_t elemPosition = 0;
+    int64_t innerElems = 1;
+    for (auto [dimSize, positionInDim] :
+         llvm::reverse(llvm::zip(vecTy.getShape(), positions))) {
+      int64_t positionVal = positionInDim.cast<IntegerAttr>().getInt();
+      elemPosition += positionVal * innerElems;
+      innerElems *= dimSize;
+    }
+
+    Attribute newAttr;
+    if (vecTy.getElementType().isIntOrIndex()) {
+      auto values = to_vector(dense.getValues<APInt>());
+      newAttr = IntegerAttr::get(extractOp.getType(), values[elemPosition]);
+    } else if (vecTy.getElementType().isa<FloatType>()) {
+      auto values = to_vector(dense.getValues<APFloat>());
+      newAttr = FloatAttr::get(extractOp.getType(), values[elemPosition]);
+    }
+    assert(newAttr && "Unhandled case");
+
+    rewriter.replaceOpWithNewOp<arith::ConstantOp>(extractOp, newAttr);
+    return success();
+  }
+};
+
 } // namespace
 
 void ExtractOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                             MLIRContext *context) {
-  results.add<ExtractOpConstantFolder, ExtractOpFromBroadcast>(context);
+  results.add<ExtractOpSplatConstantFolder, ExtractOpScalarVectorConstantFolder,
+              ExtractOpFromBroadcast>(context);
 }
 
 static void populateFromInt64AttrArray(ArrayAttr arrayAttr,
                                        SmallVectorImpl<int64_t> &results) {
   for (auto attr : arrayAttr)
     results.push_back(attr.cast<IntegerAttr>().getInt());
 }
 

mlir/test/Dialect/Vector/canonicalize.mlir

func.func @insert_extract_to_broadcast(%arg0 : vector<1x1x4xf32>,
   %arg1 : vector<4xf32>) -> (vector<4xf32>, vector<1x1x4xf32>) {
   %0 = vector.extract %arg0[0, 0] : vector<1x1x4xf32>
   %1 = vector.insert %arg1, %arg0 [0, 0] : vector<4xf32> into vector<1x1x4xf32>
   return %0, %1 : vector<4xf32>, vector<1x1x4xf32>
 }
 
 // -----
 
-// CHECK-LABEL: extract_constant
-//       CHECK-DAG: %[[CST1:.*]] = arith.constant 1 : i32
-//       CHECK-DAG: %[[CST0:.*]] = arith.constant dense<2.000000e+00> : vector<7xf32>
-//       CHECK: return %[[CST0]], %[[CST1]] : vector<7xf32>, i32
-func.func @extract_constant() -> (vector<7xf32>, i32) {
+// CHECK-LABEL: func.func @extract_splat_constant
+//   CHECK-DAG:   %[[CST1:.*]] = arith.constant 1 : i32
+//   CHECK-DAG:   %[[CST0:.*]] = arith.constant dense<2.000000e+00> : vector<7xf32>
+//  CHECK-NEXT:   return %[[CST0]], %[[CST1]] : vector<7xf32>, i32
+func.func @extract_splat_constant() -> (vector<7xf32>, i32) {
   %cst = arith.constant dense<2.000000e+00> : vector<29x7xf32>
   %cst_1 = arith.constant dense<1> : vector<4x37x9xi32>
   %0 = vector.extract %cst[2] : vector<29x7xf32>
   %1 = vector.extract %cst_1[1, 4, 5] : vector<4x37x9xi32>
   return %0, %1 : vector<7xf32>, i32
 }
 
 // -----
 
+// CHECK-LABEL: func.func @extract_1d_constant
+//   CHECK-DAG: %[[I32CST:.*]] = arith.constant 3 : i32
+//   CHECK-DAG: %[[IDXCST:.*]] = arith.constant 1 : index
+//   CHECK-DAG: %[[F32CST:.*]] = arith.constant 2.000000e+00 : f32
+//  CHECK-NEXT: return %[[I32CST]], %[[IDXCST]], %[[F32CST]] : i32, index, f32
+func.func @extract_1d_constant() -> (i32, index, f32) {
+  %icst = arith.constant dense<[1, 2, 3, 4]> : vector<4xi32>
+  %e = vector.extract %icst[2] : vector<4xi32>
+  %idx_cst = arith.constant dense<[0, 1, 2]> : vector<3xindex>
+  %f = vector.extract %idx_cst[1] : vector<3xindex>
+  %fcst = arith.constant dense<[2.000000e+00, 3.000000e+00, 4.000000e+00]> : vector<3xf32>
+  %g = vector.extract %fcst[0] : vector<3xf32>
+  return %e, %f, %g : i32, index, f32
+}
+
+// -----
+
+// CHECK-LABEL: func.func @extract_2d_constant
+//   CHECK-DAG: %[[ACST:.*]] = arith.constant 0 : i32
+//   CHECK-DAG: %[[BCST:.*]] = arith.constant 2 : i32
+//   CHECK-DAG: %[[CCST:.*]] = arith.constant 3 : i32
+//   CHECK-DAG: %[[DCST:.*]] = arith.constant 5 : i32
+//  CHECK-NEXT: return %[[ACST]], %[[BCST]], %[[CCST]], %[[DCST]] : i32, i32, i32, i32
+func.func @extract_2d_constant() -> (i32, i32, i32, i32) {
+  %cst = arith.constant dense<[[0, 1, 2], [3, 4, 5]]> : vector<2x3xi32>
+  %a = vector.extract %cst[0, 0] : vector<2x3xi32>
+  %b = vector.extract %cst[0, 2] : vector<2x3xi32>
+  %c = vector.extract %cst[1, 0] : vector<2x3xi32>
+  %d = vector.extract %cst[1, 2] : vector<2x3xi32>
+  return %a, %b, %c, %d : i32, i32, i32, i32
+}
+
+// -----
+
+// CHECK-LABEL: func.func @extract_3d_constant
+//   CHECK-DAG: %[[ACST:.*]] = arith.constant 0 : i32
+//   CHECK-DAG: %[[BCST:.*]] = arith.constant 1 : i32
+//   CHECK-DAG: %[[CCST:.*]] = arith.constant 9 : i32
+//   CHECK-DAG: %[[DCST:.*]] = arith.constant 10 : i32
+//  CHECK-NEXT: return %[[ACST]], %[[BCST]], %[[CCST]], %[[DCST]] : i32, i32, i32, i32
+func.func @extract_3d_constant() -> (i32, i32, i32, i32) {
+  %cst = arith.constant dense<[[[0, 1], [2, 3], [4, 5]], [[6, 7], [8, 9], [10, 11]]]> : vector<2x3x2xi32>
+  %a = vector.extract %cst[0, 0, 0] : vector<2x3x2xi32>
+  %b = vector.extract %cst[0, 0, 1] : vector<2x3x2xi32>
+  %c = vector.extract %cst[1, 1, 1] : vector<2x3x2xi32>
+  %d = vector.extract %cst[1, 2, 0] : vector<2x3x2xi32>
+  return %a, %b, %c, %d : i32, i32, i32, i32
+}
+
+// -----
+
 // CHECK-LABEL: extract_extract_strided
 //  CHECK-SAME: %[[A:.*]]: vector<32x16x4xf16>
 //       CHECK: %[[V:.*]] = vector.extract %[[A]][9, 7] : vector<32x16x4xf16>
 //       CHECK: return %[[V]] : vector<4xf16>
 func.func @extract_extract_strided(%arg0: vector<32x16x4xf16>) -> vector<4xf16> {
  %1 = vector.extract_strided_slice %arg0
   {offsets = [7, 3], sizes = [10, 8], strides = [1, 1]} :
   vector<32x16x4xf16> to vector<10x8x4xf16>