[clang] [CIR] Implement folder for VecShuffleDynamicOp (PR #142315)

[Amr Hesham via cfe-commits]

https://github.com/AmrDeveloper created https://github.com/llvm/llvm-project/pull/142315

This change adds a folder for the VecShuffleDynamicOp

Issue https://github.com/llvm/llvm-project/issues/136487

>From a2eaa7f28e912f6aa0df909995ec2daf665fbf4d Mon Sep 17 00:00:00 2001
From: AmrDeveloper <amr96 at programmer.net>
Date: Sun, 1 Jun 2025 19:07:56 +0200
Subject: [PATCH] [CIR] Implement folder for VecShuffleDynamicOp

---
 clang/include/clang/CIR/Dialect/IR/CIROps.td  |  1 +
 clang/lib/CIR/Dialect/IR/CIRDialect.cpp       | 32 +++++++++++++++++++
 .../Dialect/Transforms/CIRCanonicalize.cpp    |  6 ++--
 .../vector-shuffle-dynamic-fold.cir           | 18 +++++++++++
 4 files changed, 54 insertions(+), 3 deletions(-)
 create mode 100644 clang/test/CIR/Transforms/vector-shuffle-dynamic-fold.cir

diff --git a/clang/include/clang/CIR/Dialect/IR/CIROps.td b/clang/include/clang/CIR/Dialect/IR/CIROps.td
index 07851610a2abd..4442c3a5607ae 100644
--- a/clang/include/clang/CIR/Dialect/IR/CIROps.td
+++ b/clang/include/clang/CIR/Dialect/IR/CIROps.td
@@ -2188,6 +2188,7 @@ def VecShuffleDynamicOp : CIR_Op<"vec.shuffle.dynamic",
   }];
 
   let hasVerifier = 1;
+  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 36f050de9f8bb..90b32950e4774 100644
--- a/clang/lib/CIR/Dialect/IR/CIRDialect.cpp
+++ b/clang/lib/CIR/Dialect/IR/CIRDialect.cpp
@@ -1579,6 +1579,38 @@ OpFoldResult cir::VecExtractOp::fold(FoldAdaptor adaptor) {
 // VecShuffleDynamicOp
 //===----------------------------------------------------------------------===//
 
+OpFoldResult cir::VecShuffleDynamicOp::fold(FoldAdaptor adaptor) {
+  mlir::Attribute vec = adaptor.getVec();
+  mlir::Attribute indices = adaptor.getIndices();
+  if (mlir::isa_and_nonnull<cir::ConstVectorAttr>(vec) &&
+      mlir::isa_and_nonnull<cir::ConstVectorAttr>(indices)) {
+    auto vecAttr = mlir::cast<cir::ConstVectorAttr>(vec);
+    auto indicesAttr = mlir::cast<cir::ConstVectorAttr>(indices);
+    auto vecTy = cast<cir::VectorType>(vecAttr.getType());
+
+    mlir::ArrayAttr vecElts = vecAttr.getElts();
+    mlir::ArrayAttr indicesElts = indicesAttr.getElts();
+
+    const uint64_t numElements = vecElts.size();
+
+    SmallVector<mlir::Attribute, 16> elements;
+    elements.reserve(numElements);
+
+    const uint64_t maskBits = llvm::NextPowerOf2(numElements - 1) - 1;
+    for (uint64_t i = 0; i < numElements; i++) {
+      cir::IntAttr idxAttr = mlir::cast<cir::IntAttr>(indicesElts[i]);
+      uint64_t idxValue = idxAttr.getUInt();
+      uint64_t newIdx = idxValue & maskBits;
+      elements.push_back(vecElts[newIdx]);
+    }
+
+    return cir::ConstVectorAttr::get(
+        vecTy, mlir::ArrayAttr::get(getContext(), elements));
+  }
+
+  return {};
+}
+
 LogicalResult cir::VecShuffleDynamicOp::verify() {
   // The number of elements in the two input vectors must match.
   if (getVec().getType().getSize() !=
diff --git a/clang/lib/CIR/Dialect/Transforms/CIRCanonicalize.cpp b/clang/lib/CIR/Dialect/Transforms/CIRCanonicalize.cpp
index fb000adee04c6..7d03e374c27e8 100644
--- a/clang/lib/CIR/Dialect/Transforms/CIRCanonicalize.cpp
+++ b/clang/lib/CIR/Dialect/Transforms/CIRCanonicalize.cpp
@@ -138,10 +138,10 @@ void CIRCanonicalizePass::runOnOperation() {
     assert(!cir::MissingFeatures::complexRealOp());
     assert(!cir::MissingFeatures::complexImagOp());
     assert(!cir::MissingFeatures::callOp());
-    // CastOp, UnaryOp and VecExtractOp are here to perform a manual `fold` in
-    // applyOpPatternsGreedily.
+    // CastOp, UnaryOp, VecExtractOp and VecShuffleDynamicOp are here to perform
+    // a manual `fold` in applyOpPatternsGreedily.
     if (isa<BrOp, BrCondOp, CastOp, ScopeOp, SwitchOp, SelectOp, UnaryOp,
-            VecExtractOp>(op))
+            VecExtractOp, VecShuffleDynamicOp>(op))
       ops.push_back(op);
   });
 
diff --git a/clang/test/CIR/Transforms/vector-shuffle-dynamic-fold.cir b/clang/test/CIR/Transforms/vector-shuffle-dynamic-fold.cir
new file mode 100644
index 0000000000000..db5d6bff04d5b
--- /dev/null
+++ b/clang/test/CIR/Transforms/vector-shuffle-dynamic-fold.cir
@@ -0,0 +1,18 @@
+// RUN: cir-opt %s -cir-canonicalize -o - | FileCheck %s
+
+!s32i = !cir.int<s, 32>
+
+module {
+  cir.func @fold_shuffle_dynamic_vector_op_test() {
+    %alloca = cir.alloca !cir.vector<4 x !s32i>, !cir.ptr<!cir.vector<4 x !s32i>>, ["r", init]
+    %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>
+    %indices = cir.const #cir.const_vector<[#cir.int<8> : !s32i, #cir.int<7> : !s32i, #cir.int<6> : !s32i, #cir.int<5> : !s32i]> : !cir.vector<4 x !s32i>
+    %new_vec = cir.vec.shuffle.dynamic %vec : !cir.vector<4 x !s32i>, %indices : !cir.vector<4 x !s32i>
+    cir.store align(16) %new_vec, %alloca : !cir.vector<4 x !s32i>, !cir.ptr<!cir.vector<4 x !s32i>>
+    cir.return
+  }
+
+  // CHECK: %[[NEW_VEC:.*]] = cir.const #cir.const_vector<[#cir.int<1> : !s32i, #cir.int<4> : !s32i, #cir.int<3> : !s32i, #cir.int<2> : !s32i]> : !cir.vector<4 x !s32i>
+}
+
+