[Mlir-commits] [mlir] bc37077 - [mlir][Vector] Add constant folder for extractelement.

[llvmlistbot at llvm.org]

Author: jacquesguan
Date: 2022-04-02T11:10:42+08:00
New Revision: bc37077947c9d8dfb795d1448caf9f0a39d489c6

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

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

This revision adds constant folder for vector.extractelement.

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

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 f00bd0fc23999..4ffbb1a93d8c3 100644
--- a/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
+++ b/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
@@ -530,6 +530,7 @@ def Vector_ExtractElementOp :
     }
   }];
   let hasVerifier = 1;
+  let hasFolder = 1;
 }
 
 def Vector_ExtractOp :

diff  --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index 1c4ce2530ef0a..3b0dda0b14df0 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -943,6 +943,24 @@ LogicalResult vector::ExtractElementOp::verify() {
   return success();
 }
 
+OpFoldResult vector::ExtractElementOp::fold(ArrayRef<Attribute> operands) {
+  // Skip the 0-D vector here now.
+  if (operands.size() < 2)
+    return {};
+
+  Attribute src = operands[0];
+  Attribute pos = operands[1];
+  if (!src || !pos)
+    return {};
+
+  auto srcElements = src.cast<DenseElementsAttr>().getValues<Attribute>();
+
+  auto attr = pos.dyn_cast<IntegerAttr>();
+  uint64_t posIdx = attr.getInt();
+
+  return srcElements[posIdx];
+}
+
 //===----------------------------------------------------------------------===//
 // ExtractOp
 //===----------------------------------------------------------------------===//

diff  --git a/mlir/test/Dialect/Vector/canonicalize.mlir b/mlir/test/Dialect/Vector/canonicalize.mlir
index 11273fe345e5b..a083851a4bb98 100644
--- a/mlir/test/Dialect/Vector/canonicalize.mlir
+++ b/mlir/test/Dialect/Vector/canonicalize.mlir
@@ -1385,3 +1385,15 @@ func @insert_element_fold() -> vector<4xi32> {
   %1 = vector.insertelement %s, %v[%i : i32] : vector<4xi32>
   return %1 : vector<4xi32>
 }
+
+// -----
+
+// CHECK-LABEL: func @extract_element_fold
+//       CHECK:   %[[C:.+]] = arith.constant 5 : i32
+//       CHECK:   return %[[C]]
+func @extract_element_fold() -> i32 {
+  %v = arith.constant dense<[1, 3, 5, 7]> : vector<4xi32>
+  %i = arith.constant 2 : i32
+  %1 = vector.extractelement %v[%i : i32] : vector<4xi32>
+  return %1 : i32
+}