[CIR] Backport VecShuffleOp folder (#1707)

Backporting the VecShuffleOp folder

Author: AmrDeveloper

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

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

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

@@ -1096,6 +1096,40 @@ LogicalResult cir::VecTernaryOp::verify() {
 // VecShuffle
 //===----------------------------------------------------------------------===//
 
+OpFoldResult cir::VecShuffleOp::fold(FoldAdaptor adaptor) {
+  auto vec1Attr =
+      mlir::dyn_cast_if_present<cir::ConstVectorAttr>(adaptor.getVec1());
+  auto vec2Attr =
+      mlir::dyn_cast_if_present<cir::ConstVectorAttr>(adaptor.getVec2());
+  if (!vec1Attr || !vec2Attr)
+    return {};
+
+  mlir::Type vec1ElemTy =
+      mlir::cast<cir::VectorType>(vec1Attr.getType()).getElementType();
+
+  mlir::ArrayAttr vec1Elts = vec1Attr.getElts();
+  mlir::ArrayAttr vec2Elts = vec2Attr.getElts();
+  mlir::ArrayAttr indicesElts = adaptor.getIndices();
+
+  SmallVector<mlir::Attribute, 16> elements;
+  elements.reserve(indicesElts.size());
+
+  uint64_t vec1Size = vec1Elts.size();
+  for (const auto &idxAttr : indicesElts.getAsRange<cir::IntAttr>()) {
+    if (idxAttr.getSInt() == -1) {
+      elements.push_back(cir::UndefAttr::get(vec1ElemTy));
+      continue;
+    }
+
+    uint64_t idxValue = idxAttr.getUInt();
+    elements.push_back(idxValue < vec1Size ? vec1Elts[idxValue]
+                                           : vec2Elts[idxValue - vec1Size]);
+  }
+
+  return cir::ConstVectorAttr::get(
+      getType(), mlir::ArrayAttr::get(getContext(), elements));
+}
+
 LogicalResult cir::VecShuffleOp::verify() {
   // The number of elements in the indices array must match the number of
   // elements in the result type.

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, VecTernaryOp>(op))
+            VecExtractOp, VecShuffleOp, VecTernaryOp>(op))
       ops.push_back(op);
   });
 

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

@@ -0,0 +1,59 @@
+// RUN: cir-opt %s -cir-canonicalize -o - -split-input-file | FileCheck %s
+
+!s32i = !cir.int<s, 32>
+!s64i = !cir.int<s, 64>
+
+module  {
+  cir.func @fold_shuffle_vector_op_test() -> !cir.vector<!s32i x 4> {
+    %vec_1 = cir.const #cir.const_vector<[#cir.int<1> : !s32i, #cir.int<3> : !s32i, #cir.int<5> : !s32i, #cir.int<7> : !s32i]> : !cir.vector<!s32i x 4>
+    %vec_2 = cir.const #cir.const_vector<[#cir.int<2> : !s32i, #cir.int<4> : !s32i, #cir.int<6> : !s32i, #cir.int<8> : !s32i]> : !cir.vector<!s32i x 4>
+    %new_vec = cir.vec.shuffle(%vec_1, %vec_2 : !cir.vector<!s32i x 4>) [#cir.int<0> : !s64i, #cir.int<4> : !s64i,
+      #cir.int<1> : !s64i, #cir.int<5> : !s64i] : !cir.vector<!s32i x 4>
+    cir.return %new_vec : !cir.vector<!s32i x 4>
+  }
+
+  // CHECK: cir.func @fold_shuffle_vector_op_test() -> !cir.vector<!s32i x 4> {
+  // CHECK-NEXT: %[[RES:.*]] = cir.const #cir.const_vector<[#cir.int<1> : !s32i, #cir.int<2> : !s32i, #cir.int<3> : !s32i,
+  // CHECK-SAME: #cir.int<4> : !s32i]> : !cir.vector<!s32i x 4>
+  // CHECK-NEXT: cir.return %[[RES]] : !cir.vector<!s32i x 4>
+}
+
+// -----
+
+!s32i = !cir.int<s, 32>
+!s64i = !cir.int<s, 64>
+
+module  {
+  cir.func @fold_shuffle_vector_op_test() -> !cir.vector<!s32i x 6> {
+    %vec_1 = cir.const #cir.const_vector<[#cir.int<1> : !s32i, #cir.int<3> : !s32i, #cir.int<5> : !s32i, #cir.int<7> : !s32i]> : !cir.vector<!s32i x 4>
+    %vec_2 = cir.const #cir.const_vector<[#cir.int<2> : !s32i, #cir.int<4> : !s32i, #cir.int<6> : !s32i, #cir.int<8> : !s32i]> : !cir.vector<!s32i x 4>
+    %new_vec = cir.vec.shuffle(%vec_1, %vec_2 : !cir.vector<!s32i x 4>) [#cir.int<0> : !s64i, #cir.int<4> : !s64i,
+      #cir.int<1> : !s64i, #cir.int<5> : !s64i, #cir.int<2> : !s64i, #cir.int<6> : !s64i] : !cir.vector<!s32i x 6>
+    cir.return %new_vec : !cir.vector<!s32i x 6>
+  }
+
+  // CHECK: cir.func @fold_shuffle_vector_op_test() -> !cir.vector<!s32i x 6> {
+  // CHECK-NEXT: %[[RES:.*]] = cir.const #cir.const_vector<[#cir.int<1> : !s32i, #cir.int<2> : !s32i, #cir.int<3> : !s32i,
+  // CHECK-SAME: #cir.int<4> : !s32i, #cir.int<5> : !s32i, #cir.int<6> : !s32i]> : !cir.vector<!s32i x 6>
+  // CHECK-NEXT: cir.return %[[RES]] : !cir.vector<!s32i x 6>
+}
+
+// -----
+
+!s32i = !cir.int<s, 32>
+!s64i = !cir.int<s, 64>
+
+module  {
+  cir.func @fold_shuffle_vector_op_test() -> !cir.vector<!s32i x 4> {
+    %vec_1 = cir.const #cir.const_vector<[#cir.int<1> : !s32i, #cir.int<3> : !s32i, #cir.int<5> : !s32i, #cir.int<7> : !s32i]> : !cir.vector<!s32i x 4>
+    %vec_2 = cir.const #cir.const_vector<[#cir.int<2> : !s32i, #cir.int<4> : !s32i, #cir.int<6> : !s32i, #cir.int<8> : !s32i]> : !cir.vector<!s32i x 4>
+    %new_vec = cir.vec.shuffle(%vec_1, %vec_2 : !cir.vector<!s32i x 4>) [#cir.int<-1> : !s64i, #cir.int<4> : !s64i,
+      #cir.int<1> : !s64i, #cir.int<5> : !s64i] : !cir.vector<!s32i x 4>
+    cir.return %new_vec : !cir.vector<!s32i x 4>
+  }
+
+  // CHECK: cir.func @fold_shuffle_vector_op_test() -> !cir.vector<!s32i x 4> {
+  // CHECK: cir.const #cir.const_vector<[#cir.undef : !s32i, #cir.int<2> : !s32i, #cir.int<3> : !s32i,
+  // CHECK-SAME: #cir.int<4> : !s32i]> : !cir.vector<!s32i x 4>
+  // CHECK-NEXT: cir.return %[[RES]] : !cir.vector<!s32i x 4>
+}