[CIR] Backport folder implementation for VecExtractOp (llvm#1613)

Backport llvm#139304

Author: AmrDeveloper

clang/include/clang/CIR/Dialect/IR/CIROps.td

@@ -3165,6 +3165,7 @@ def VecExtractOp : CIR_Op<"vec.extract", [Pure,
   }];
 
   let llvmOp = "ExtractElementOp";
+  let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//

clang/lib/CIR/Dialect/IR/CIRDialect.cpp

@@ -1098,6 +1098,29 @@ LogicalResult cir::VecShuffleDynamicOp::verify() {
   return success();
 }
 
+//===----------------------------------------------------------------------===//
+// VecExtractOp
+//===----------------------------------------------------------------------===//
+
+OpFoldResult cir::VecExtractOp::fold(FoldAdaptor adaptor) {
+  const auto vectorAttr =
+      llvm::dyn_cast_if_present<cir::ConstVectorAttr>(adaptor.getVec());
+  if (!vectorAttr)
+    return {};
+
+  const auto indexAttr =
+      llvm::dyn_cast_if_present<cir::IntAttr>(adaptor.getIndex());
+  if (!indexAttr)
+    return {};
+
+  const mlir::ArrayAttr elements = vectorAttr.getElts();
+  const uint64_t index = indexAttr.getUInt();
+  if (index >= elements.size())
+    return {};
+
+  return elements[index];
+}
+
 //===----------------------------------------------------------------------===//
 // ReturnOp
 //===----------------------------------------------------------------------===//

clang/lib/CIR/Dialect/Transforms/CIRCanonicalize.cpp

@@ -79,9 +79,9 @@ struct RemoveEmptyScope : public OpRewritePattern<ScopeOp> {
 struct RemoveEmptySwitch : public OpRewritePattern<SwitchOp> {
   using OpRewritePattern<SwitchOp>::OpRewritePattern;
 
-  LogicalResult matchAndRewrite(SwitchOp op, PatternRewriter &rewriter) const final {
-    if (!(op.getBody().empty() ||
-                   isa<YieldOp>(op.getBody().front().front())))
+  LogicalResult matchAndRewrite(SwitchOp op,
+                                PatternRewriter &rewriter) const final {
+    if (!(op.getBody().empty() || isa<YieldOp>(op.getBody().front().front())))
       return failure();
 
     rewriter.eraseOp(op);
@@ -92,7 +92,8 @@ struct RemoveEmptySwitch : public OpRewritePattern<SwitchOp> {
 struct RemoveTrivialTry : public OpRewritePattern<TryOp> {
   using OpRewritePattern<TryOp>::OpRewritePattern;
 
-  LogicalResult matchAndRewrite(TryOp op, PatternRewriter &rewriter) const final {
+  LogicalResult matchAndRewrite(TryOp op,
+                                PatternRewriter &rewriter) const final {
     // FIXME: also check all catch regions are empty
     // return success(op.getTryRegion().hasOneBlock());
     return mlir::failure();
@@ -116,7 +117,8 @@ struct RemoveTrivialTry : public OpRewritePattern<TryOp> {
 struct SimplifyCallOp : public OpRewritePattern<CallOp> {
   using OpRewritePattern<CallOp>::OpRewritePattern;
 
-  LogicalResult matchAndRewrite(CallOp op, PatternRewriter &rewriter) const final {
+  LogicalResult matchAndRewrite(CallOp op,
+                                PatternRewriter &rewriter) const final {
     // Applicable to cir.call exception ... clean { cir.yield }
     mlir::Region *r = &op.getCleanup();
     if (r->empty() || !r->hasOneBlock())
@@ -174,10 +176,11 @@ void CIRCanonicalizePass::runOnOperation() {
   // Collect operations to apply patterns.
   llvm::SmallVector<Operation *, 16> ops;
   getOperation()->walk([&](Operation *op) {
-    // CastOp and UnaryOp are here to perform a manual `fold` in
+    // CastOp, UnaryOp and VecExtractOp are here to perform a manual `fold` in
     // applyOpPatternsGreedily.
     if (isa<BrOp, BrCondOp, ScopeOp, SwitchOp, CastOp, TryOp, UnaryOp, SelectOp,
-            ComplexCreateOp, ComplexRealOp, ComplexImagOp, CallOp>(op))
+            ComplexCreateOp, ComplexRealOp, ComplexImagOp, CallOp,
+            VecExtractOp>(op))
       ops.push_back(op);
   });
 

clang/test/CIR/Transforms/vector-extract-fold.cir

@@ -0,0 +1,33 @@
+// RUN: cir-opt %s -cir-canonicalize -o - | FileCheck %s
+
+!s32i = !cir.int<s, 32>
+
+module {
+  cir.func @fold_extract_vector_op_test() {
+    %init = cir.alloca !s32i, !cir.ptr<!s32i>, ["e", init]
+    %const_vec = cir.const #cir.const_vector<[#cir.int<1> : !s32i, #cir.int<2> : !s32i, #cir.int<3> : !s32i, #cir.int<4> : !s32i]> : !cir.vector<!s32i x 4>
+    %index = cir.const #cir.int<1> : !s32i
+    %ele = cir.vec.extract %const_vec[%index : !s32i] : !cir.vector<!s32i x 4>
+    cir.store %ele, %init : !s32i, !cir.ptr<!s32i>
+    cir.return
+  }
+
+  // CHECK: %[[INIT:.*]] = cir.alloca !s32i, !cir.ptr<!s32i>, ["e", init]
+  // CHECK: %[[VALUE:.*]] = cir.const #cir.int<2> : !s32i
+  // CHECK: cir.store %[[VALUE]], %[[INIT]] : !s32i, !cir.ptr<!s32i>
+
+  cir.func @fold_extract_vector_op_index_out_of_bounds_test() {
+    %init = cir.alloca !s32i, !cir.ptr<!s32i>, ["e", init]
+    %const_vec = cir.const #cir.const_vector<[#cir.int<1> : !s32i, #cir.int<2> : !s32i, #cir.int<3> : !s32i, #cir.int<4> : !s32i]> : !cir.vector<!s32i x 4>
+    %index = cir.const #cir.int<9> : !s32i
+    %ele = cir.vec.extract %const_vec[%index : !s32i] : !cir.vector<!s32i x 4>
+    cir.store %ele, %init : !s32i, !cir.ptr<!s32i>
+    cir.return
+  }
+
+  // CHECK: %[[INIT:.*]] = cir.alloca !s32i, !cir.ptr<!s32i>, ["e", init]
+  // CHECK: %[[CONST_VEC:.*]] = cir.const #cir.const_vector<[#cir.int<1> : !s32i, #cir.int<2> : !s32i, #cir.int<3> : !s32i, #cir.int<4> : !s32i]> : !cir.vector<!s32i x 4>
+  // CHECK: %[[INDEX:.*]] = cir.const #cir.int<9> : !s32i
+  // CHECK: %[[ELE:.*]] = cir.vec.extract %[[CONST_VEC]][%[[INDEX]] : !s32i] : !cir.vector<!s32i x 4>
+  // CHECK: cir.store %[[ELE]], %[[INIT]] : !s32i, !cir.ptr<!s32i>
+}