From d9380ec63707be2cf0b0010220e4a34d2dbd107b Mon Sep 17 00:00:00 2001
From: Amr Hesham <amr96@programmer.net>
Date: Tue, 13 May 2025 20:01:48 +0200
Subject: [PATCH] [CIR] Implement folder for VecExtractOp (#139304)

This change adds a folder for the VecExtractOp

Issue https://github.com/llvm/llvm-project/issues/136487
---
 clang/include/clang/CIR/Dialect/IR/CIROps.td  |  2 ++
 clang/lib/CIR/Dialect/IR/CIRDialect.cpp       | 23 ++++++++++++
 .../Dialect/Transforms/CIRCanonicalize.cpp    |  5 +--
 .../CIR/Transforms/vector-extract-fold.cir    | 35 +++++++++++++++++++
 4 files changed, 63 insertions(+), 2 deletions(-)
 create mode 100644 clang/test/CIR/Transforms/vector-extract-fold.cir

diff --git a/clang/include/clang/CIR/Dialect/IR/CIROps.td b/clang/include/clang/CIR/Dialect/IR/CIROps.td
index e8020b0ad803..251e78aba254 100644
--- a/clang/include/clang/CIR/Dialect/IR/CIROps.td
+++ b/clang/include/clang/CIR/Dialect/IR/CIROps.td
@@ -2031,6 +2031,8 @@ def VecExtractOp : CIR_Op<"vec.extract", [Pure,
   let assemblyFormat = [{
     $vec `[` $index `:` type($index) `]` attr-dict `:` qualified(type($vec))
   }];
+
+  let hasFolder = 1;
 }
 
 #endif // CLANG_CIR_DIALECT_IR_CIROPS_TD
diff --git a/clang/lib/CIR/Dialect/IR/CIRDialect.cpp b/clang/lib/CIR/Dialect/IR/CIRDialect.cpp
index b131edaf403e..abb0928b377c 100644
--- a/clang/lib/CIR/Dialect/IR/CIRDialect.cpp
+++ b/clang/lib/CIR/Dialect/IR/CIRDialect.cpp
@@ -1395,6 +1395,29 @@ LogicalResult cir::VecCreateOp::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];
+}
+
 //===----------------------------------------------------------------------===//
 // TableGen'd op method definitions
 //===----------------------------------------------------------------------===//
diff --git a/clang/lib/CIR/Dialect/Transforms/CIRCanonicalize.cpp b/clang/lib/CIR/Dialect/Transforms/CIRCanonicalize.cpp
index 3b4c7bc61313..798bc0dab938 100644
--- a/clang/lib/CIR/Dialect/Transforms/CIRCanonicalize.cpp
+++ b/clang/lib/CIR/Dialect/Transforms/CIRCanonicalize.cpp
@@ -125,9 +125,10 @@ void CIRCanonicalizePass::runOnOperation() {
     assert(!cir::MissingFeatures::complexRealOp());
     assert(!cir::MissingFeatures::complexImagOp());
     assert(!cir::MissingFeatures::callOp());
-    // 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, CastOp, ScopeOp, SelectOp, UnaryOp>(op))
+    if (isa<BrOp, BrCondOp, CastOp, ScopeOp, SelectOp, UnaryOp, VecExtractOp>(
+            op))
       ops.push_back(op);
   });
 
diff --git a/clang/test/CIR/Transforms/vector-extract-fold.cir b/clang/test/CIR/Transforms/vector-extract-fold.cir
new file mode 100644
index 000000000000..14fed50454b8
--- /dev/null
+++ b/clang/test/CIR/Transforms/vector-extract-fold.cir
@@ -0,0 +1,35 @@
+// 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<4 x !s32i>
+    %index = cir.const #cir.int<1> : !s32i
+    %ele = cir.vec.extract %const_vec[%index : !s32i] : !cir.vector<4 x !s32i>
+    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<4 x !s32i>
+    %index = cir.const #cir.int<9> : !s32i
+    %ele = cir.vec.extract %const_vec[%index : !s32i] : !cir.vector<4 x !s32i>
+    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<4 x !s32i>
+  // CHECK: %[[INDEX:.*]] = cir.const #cir.int<9> : !s32i
+  // CHECK: %[[ELE:.*]] = cir.vec.extract %[[CONST_VEC]][%[[INDEX]] : !s32i] : !cir.vector<4 x !s32i>
+  // CHECK: cir.store %[[ELE]], %[[INIT]] : !s32i, !cir.ptr<!s32i>
+}
+
+