[CIR] Backport VecShuffleDynamicOp folder (#1708)

Backporting the VecShuffleDynamicOp folder

Author: AmrDeveloper

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

@@ -3320,7 +3320,9 @@ def VecShuffleDynamicOp : CIR_Op<"vec.shuffle.dynamic",
     $vec `:` qualified(type($vec)) `,` $indices `:` qualified(type($indices))
     attr-dict
   }];
+
   let hasVerifier = 1;
+  let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//

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

@@ -1165,6 +1165,34 @@ LogicalResult cir::VecShuffleOp::verify() {
 // VecShuffleDynamic
 //===----------------------------------------------------------------------===//
 
+OpFoldResult cir::VecShuffleDynamicOp::fold(FoldAdaptor adaptor) {
+  auto vecAttr =
+      mlir::dyn_cast_if_present<cir::ConstVectorAttr>(adaptor.getVec());
+  auto indicesAttr =
+      mlir::dyn_cast_if_present<cir::ConstVectorAttr>(adaptor.getIndices());
+  if (!vecAttr || !indicesAttr)
+    return {};
+
+  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 (const auto &idxAttr : indicesElts.getAsRange<cir::IntAttr>()) {
+    uint64_t idxValue = idxAttr.getUInt();
+    uint64_t newIdx = idxValue & maskBits;
+    elements.push_back(vecElts[newIdx]);
+  }
+
+  auto vecTy = mlir::cast<cir::VectorType>(vecAttr.getType());
+  return cir::ConstVectorAttr::get(
+      vecTy, mlir::ArrayAttr::get(getContext(), elements));
+}
+
 LogicalResult cir::VecShuffleDynamicOp::verify() {
   // The number of elements in the two input vectors must match.
   if (getVec().getType().getSize() !=

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

@@ -180,7 +180,7 @@ void CIRCanonicalizePass::runOnOperation() {
     // applyOpPatternsGreedily.
     if (isa<BrOp, BrCondOp, ScopeOp, SwitchOp, CastOp, TryOp, UnaryOp, SelectOp,
             ComplexCreateOp, ComplexRealOp, ComplexImagOp, CallOp, VecCreateOp,
-            VecExtractOp, VecShuffleOp, VecTernaryOp>(op))
+            VecExtractOp, VecShuffleOp, VecShuffleDynamicOp, VecTernaryOp>(op))
       ops.push_back(op);
   });
 

clang/test/CIR/Transforms/vector-shuffle-dynamic-fold.cir

@@ -0,0 +1,30 @@
+// RUN: cir-opt %s -cir-canonicalize -o - | FileCheck %s
+
+!s32i = !cir.int<s, 32>
+
+module {
+  cir.func @fold_shuffle_dynamic_vector_op_test() -> !cir.vector<!s32i x 4> {
+    %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>
+    %indices = cir.const #cir.const_vector<[#cir.int<8> : !s32i, #cir.int<7> : !s32i, #cir.int<6> : !s32i, #cir.int<5> : !s32i]> : !cir.vector<!s32i x 4>
+    %new_vec = cir.vec.shuffle.dynamic %vec : !cir.vector<!s32i x 4>, %indices : !cir.vector<!s32i x 4>
+    cir.return %new_vec : !cir.vector<!s32i x 4>
+  }
+
+  // Masking indices [8, 7, 6, 5] AND 3 = [0, 3, 2, 1]
+  // CHECK: cir.func @fold_shuffle_dynamic_vector_op_test() -> !cir.vector<!s32i x 4> {
+  // CHECK-NEXT: %[[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<!s32i x 4>
+  // CHECK-NEXT: cir.return %[[NEW_VEC:.*]] : !cir.vector<!s32i x 4>
+
+  cir.func @fold_shuffle_dynamic_vector_op_test_2() -> !cir.vector<!s32i x 4> {
+    %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>
+    %indices = cir.const #cir.const_vector<[#cir.int<3> : !s32i, #cir.int<2> : !s32i, #cir.int<1> : !s32i, #cir.int<0> : !s32i]> : !cir.vector<!s32i x 4>
+    %new_vec = cir.vec.shuffle.dynamic %vec : !cir.vector<!s32i x 4>, %indices : !cir.vector<!s32i x 4>
+    cir.return %new_vec : !cir.vector<!s32i x 4>
+  }
+
+  // Masking indices [3, 2, 1, 0] AND 3 = [3, 2, 1, 0]
+  // CHECK: cir.func @fold_shuffle_dynamic_vector_op_test_2() -> !cir.vector<!s32i x 4> {
+  // CHECK-NEXT: %[[NEW_VEC:.*]] = cir.const #cir.const_vector<[#cir.int<4> : !s32i, #cir.int<3> : !s32i, #cir.int<2> : !s32i, #cir.int<1> : !s32i]> : !cir.vector<!s32i x 4>
+  // CHECK-NEXT: cir.return %[[NEW_VEC:.*]] : !cir.vector<!s32i x 4>
+}
+