[Mlir-commits] [mlir] 2628236 - [mlir][Vector] Add constant folder for insertelement.

[llvmlistbot at llvm.org]

Author: jacquesguan
Date: 2022-04-02T10:20:19+08:00
New Revision: 262823612d616134b0fe225ab3e848f21d6e31ff

URL: https://github.com/llvm/llvm-project/commit/262823612d616134b0fe225ab3e848f21d6e31ff
DIFF: https://github.com/llvm/llvm-project/commit/262823612d616134b0fe225ab3e848f21d6e31ff.diff

LOG: [mlir][Vector] Add constant folder for insertelement.

This revision adds constant folder for vector.insertelement.

Differential Revision: https://reviews.llvm.org/D122721

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
    mlir/lib/Dialect/Vector/IR/VectorOps.cpp
    mlir/test/Dialect/Vector/canonicalize.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td b/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
index 638e69ed7eef8..f00bd0fc23999 100644
--- a/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
+++ b/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
@@ -713,6 +713,7 @@ def Vector_InsertElementOp :
     }
   }];
   let hasVerifier = 1;
+  let hasFolder = 1;
 }
 
 def Vector_InsertOp :

diff  --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index c7d60e5e3aab6..1c4ce2530ef0a 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -1839,6 +1839,29 @@ LogicalResult InsertElementOp::verify() {
   return success();
 }
 
+OpFoldResult vector::InsertElementOp::fold(ArrayRef<Attribute> operands) {
+  // Skip the 0-D vector here.
+  if (operands.size() < 3)
+    return {};
+
+  Attribute src = operands[0];
+  Attribute dst = operands[1];
+  Attribute pos = operands[2];
+  if (!src || !dst || !pos)
+    return {};
+
+  auto dstElements = dst.cast<DenseElementsAttr>().getValues<Attribute>();
+
+  SmallVector<Attribute> results(dstElements);
+
+  auto attr = pos.dyn_cast<IntegerAttr>();
+  uint64_t posIdx = attr.getInt();
+
+  results[posIdx] = src;
+
+  return DenseElementsAttr::get(getDestVectorType(), results);
+}
+
 //===----------------------------------------------------------------------===//
 // InsertOp
 //===----------------------------------------------------------------------===//

diff  --git a/mlir/test/Dialect/Vector/canonicalize.mlir b/mlir/test/Dialect/Vector/canonicalize.mlir
index 13022c29cd4e8..11273fe345e5b 100644
--- a/mlir/test/Dialect/Vector/canonicalize.mlir
+++ b/mlir/test/Dialect/Vector/canonicalize.mlir
@@ -1372,3 +1372,16 @@ func @transpose_scalar_broadcast2(%value: f32) -> vector<1x8xf32> {
   %t = vector.transpose %bcast, [1, 0] : vector<8x1xf32> to vector<1x8xf32>
   return %t : vector<1x8xf32>
 }
+
+// -----
+
+// CHECK-LABEL: func @insert_element_fold
+//       CHECK:   %[[V:.+]] = arith.constant dense<[0, 1, 7, 3]> : vector<4xi32>
+//       CHECK:   return %[[V]]
+func @insert_element_fold() -> vector<4xi32> {
+  %v = arith.constant dense<[0, 1, 2, 3]> : vector<4xi32>
+  %s = arith.constant 7 : i32
+  %i = arith.constant 2 : i32
+  %1 = vector.insertelement %s, %v[%i : i32] : vector<4xi32>
+  return %1 : vector<4xi32>
+}