update function name and add canonicalize test and regression test.

Author: linuxlonelyeagle

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

@@ -1979,33 +1979,27 @@ static Value foldScalarExtractFromFromElements(ExtractOp extractOp) {
 
 // If the dynamic operands of `extractOp` or `insertOp` is result of
 // `constantOp`, then fold it.
-template <typename T>
-static LogicalResult foldConstantOp(T op, SmallVectorImpl<Value> &operands) {
+template <typename OpType, typename AdaptorType>
+static Value extractInsertFoldConstantOp(OpType op, AdaptorType adaptor,
+                                         SmallVectorImpl<Value> &operands) {
   auto staticPosition = op.getStaticPosition().vec();
   OperandRange dynamicPosition = op.getDynamicPosition();
-
+  ArrayRef<Attribute> dynamicPositionAttr = adaptor.getDynamicPosition();
   // If the dynamic operands is empty, it is returned directly.
   if (!dynamicPosition.size())
-    return failure();
+    return {};
   unsigned index = 0;
 
-  // `opChange` is a flog. If it is true, it means to update `op` in place.
+  // `opChange` is a flag. If it is true, it means to update `op` in place.
   bool opChange = false;
   for (unsigned i = 0, e = staticPosition.size(); i < e; ++i) {
     if (!ShapedType::isDynamic(staticPosition[i]))
       continue;
+    Attribute positionAttr = dynamicPositionAttr[index];
     Value position = dynamicPosition[index++];
-
-    // If it is a block parameter, proceed to the next iteration.
-    if (!position.getDefiningOp()) {
-      operands.push_back(position);
-      continue;
-    }
-
-    APInt pos;
-    if (matchPattern(position, m_ConstantInt(&pos))) {
+    if (auto attr = mlir::dyn_cast_if_present<IntegerAttr>(positionAttr)) {
+      staticPosition[i] = attr.getInt();
       opChange = true;
-      staticPosition[i] = pos.getSExtValue();
       continue;
     }
     operands.push_back(position);
@@ -2014,12 +2008,12 @@ static LogicalResult foldConstantOp(T op, SmallVectorImpl<Value> &operands) {
   if (opChange) {
     op.setStaticPosition(staticPosition);
     op.getOperation()->setOperands(operands);
-    return success();
+    return op.getResult();
   }
-  return failure();
+  return {};
 }
 
-OpFoldResult ExtractOp::fold(FoldAdaptor) {
+OpFoldResult ExtractOp::fold(FoldAdaptor adaptor) {
   // Fold "vector.extract %v[] : vector<2x2xf32> from vector<2x2xf32>" to %v.
   // Note: Do not fold "vector.extract %v[] : f32 from vector<f32>" (type
   // mismatch).
@@ -2042,8 +2036,8 @@ OpFoldResult ExtractOp::fold(FoldAdaptor) {
   if (auto val = foldScalarExtractFromFromElements(*this))
     return val;
   SmallVector<Value> operands = {getVector()};
-  if (succeeded(foldConstantOp(*this, operands)))
-    return getResult();
+  if (auto val = extractInsertFoldConstantOp(*this, adaptor, operands))
+    return val;
   return OpFoldResult();
 }
 
@@ -3074,8 +3068,8 @@ OpFoldResult vector::InsertOp::fold(FoldAdaptor adaptor) {
   if (getNumIndices() == 0 && getSourceType() == getType())
     return getSource();
   SmallVector<Value> operands = {getSource(), getDest()};
-  if (succeeded(foldConstantOp(*this, operands)))
-    return getResult();
+  if (auto val = extractInsertFoldConstantOp(*this, adaptor, operands))
+    return val;
   return {};
 }
 

mlir/test/Conversion/VectorToLLVM/vector-to-llvm.mlir

@@ -4115,39 +4115,3 @@ func.func @step_scalable() -> vector<[4]xindex> {
   %0 = vector.step : vector<[4]xindex>
   return %0 : vector<[4]xindex>
 }
-
-// -----
-
-// CHECK-LABEL: @extract_arith_constnt
-func.func @extract_arith_constnt() -> i32 {
-  %v = arith.constant dense<0> : vector<32x1xi32>
-  %c_0 = arith.constant 0 : index
-  %elem = vector.extract %v[%c_0, %c_0] : i32 from vector<32x1xi32>
-  return %elem : i32
-}
-
-// CHECK: %[[VAL_0:.*]] = arith.constant dense<0> : vector<32x1xi32>
-// CHECK: %[[VAL_1:.*]] = builtin.unrealized_conversion_cast %[[VAL_0]] : vector<32x1xi32> to !llvm.array<32 x vector<1xi32>>
-// CHECK: %[[VAL_2:.*]] = llvm.extractvalue %[[VAL_1]][0] : !llvm.array<32 x vector<1xi32>>
-// CHECK: %[[VAL_3:.*]] = llvm.mlir.constant(0 : i64) : i64
-// CHECK: %{{.*}} = llvm.extractelement %[[VAL_2]]{{\[}}%[[VAL_3]] : i64] : vector<1xi32>
-
-// -----
-
-// CHECK-LABEL: @insert_arith_constnt()
-
-func.func @insert_arith_constnt() -> vector<32x1xi32> {
-  %v = arith.constant dense<0> : vector<32x1xi32>
-  %c_0 = arith.constant 0 : index
-  %c_1 = arith.constant 1 : i32
-  %v_1 = vector.insert %c_1, %v[%c_0, %c_0] : i32 into vector<32x1xi32>
-  return %v_1 : vector<32x1xi32>
-}
-    
-// CHECK: %[[VAL_0:.*]] = arith.constant dense<0> : vector<32x1xi32>
-// CHECK: %[[VAL_1:.*]] = builtin.unrealized_conversion_cast %[[VAL_0]] : vector<32x1xi32> to !llvm.array<32 x vector<1xi32>>
-// CHECK: %[[VAL_2:.*]] = arith.constant 1 : i32
-// CHECK: %[[VAL_3:.*]] = llvm.extractvalue %[[VAL_1]][0] : !llvm.array<32 x vector<1xi32>>
-// CHECK: %[[VAL_4:.*]] = llvm.mlir.constant(0 : i64) : i64
-// CHECK: %[[VAL_5:.*]] = llvm.insertelement %[[VAL_2]], %[[VAL_3]]{{\[}}%[[VAL_4]] : i64] : vector<1xi32>
-// CHECK: %{{.*}} = llvm.insertvalue %[[VAL_5]], %[[VAL_1]][0] : !llvm.array<32 x vector<1xi32>>

mlir/test/Dialect/Vector/canonicalize.mlir

@@ -2979,3 +2979,47 @@ func.func @contiguous_scatter_step(%base: memref<?xf32>,
     memref<?xf32>, vector<16xindex>, vector<16xi1>, vector<16xf32>
   return
 }
+
+// -----
+
+// CHECK-LABEL: func @extract_arith_constnt
+
+func.func @extract_arith_constnt() -> i32 {
+  %c1_i32 = arith.constant 1 : i32
+  return %c1_i32 : i32
+}
+
+// CHECK: %[[VAL_0:.*]] = arith.constant 1 : i32
+// CHECK: return %[[VAL_0]] : i32
+
+// -----
+
+// CHECK-LABEL: func @insert_arith_constnt
+
+func.func @insert_arith_constnt() -> vector<4x1xi32> {
+  %v = arith.constant dense<0> : vector<4x1xi32>
+  %c_0 = arith.constant 0 : index
+  %c_1 = arith.constant 1 : i32
+  %v_1 = vector.insert %c_1, %v[%c_0, %c_0] : i32 into vector<4x1xi32>
+  return %v_1 : vector<4x1xi32>
+}
+
+// CHECK: %[[VAL_0:.*]] = arith.constant dense<{{\[\[}}1], [0], [0], [0]]> : vector<4x1xi32>
+// CHECK: return %[[VAL_0]] : vector<4x1xi32>
+
+// -----
+
+// CHECK-LABEL: func @insert_extract_arith_constnt
+
+func.func @insert_extract_arith_constnt() -> i32 {
+  %v = arith.constant dense<0> : vector<32x1xi32>
+  %c_0 = arith.constant 0 : index
+  %c_1 = arith.constant 1 : index
+  %c_2 = arith.constant 2 : i32
+  %v_1 = vector.insert %c_2, %v[%c_1, %c_1] : i32 into vector<32x1xi32>
+  %ret = vector.extract %v_1[%c_1, %c_1] : i32 from vector<32x1xi32>
+  return %ret : i32
+}
+
+// CHECK: %[[VAL_0:.*]] = arith.constant 2 : i32
+// CHECK: return %[[VAL_0]] : i32

mlir/test/Integration/Dialect/Vector/CPU/extract-insert-fold-constant.mlir

@@ -0,0 +1,34 @@
+// RUN: mlir-opt %s -test-lower-to-llvm  | \
+// RUN: mlir-runner -e entry -entry-point-result=void  \
+// RUN:   -shared-libs=%mlir_c_runner_utils | \
+// RUN: FileCheck %s
+
+func.func @entry() {
+  %v = arith.constant dense<0> : vector<2x2xi32>
+  %c_0 = arith.constant 0 : index
+  %c_1 = arith.constant 1 : index
+  %i32_0 = arith.constant 0 : i32
+  %i32_1 = arith.constant 1 : i32
+  %i32_2 = arith.constant 2 : i32
+  %i32_3 = arith.constant 3 : i32
+  %v_1 = vector.insert %i32_0, %v[%c_0, %c_0] : i32 into vector<2x2xi32>
+  %v_2 = vector.insert %i32_1, %v_1[%c_0, %c_1] : i32 into vector<2x2xi32>
+  %v_3 = vector.insert %i32_2, %v_2[%c_1, %c_0] : i32 into vector<2x2xi32>
+  %v_4 = vector.insert %i32_3, %v_3[%c_1, %c_1] : i32 into vector<2x2xi32>
+  // CHECK: ( ( 0, 1 ), ( 2, 3 ) ) 
+  vector.print %v_4 : vector<2x2xi32>
+  %v_5 = arith.constant dense<[[0, 1], [2, 3]]> : vector<2x2xi32>
+  // CHECK: 0
+  %i32_4 = vector.extract %v_5[%c_0, %c_0] : i32 from vector<2x2xi32>
+  // CHECK: 1
+  %i32_5 = vector.extract %v_5[%c_0, %c_1] : i32 from vector<2x2xi32>
+  // CHECK: 2
+  %i32_6 = vector.extract %v_5[%c_1, %c_0] : i32 from vector<2x2xi32>
+  // CHECK: 3
+  %i32_7 = vector.extract %v_5[%c_1, %c_1] : i32 from vector<2x2xi32>
+  vector.print %i32_4 : i32
+  vector.print %i32_5 : i32
+  vector.print %i32_6 : i32
+  vector.print %i32_7 : i32
+  return
+}