[CIR] Backport VecTernaryOp folder (#1705)

Backporting the VecTernaryOp folder

Author: AmrDeveloper

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

@@ -3225,7 +3225,7 @@ def VecCmpOp : CIR_Op<"vec.cmp", [Pure, SameTypeOperands]> {
 //===----------------------------------------------------------------------===//
 
 def VecTernaryOp : CIR_Op<"vec.ternary",
-                   [Pure, AllTypesMatch<["result", "vec1", "vec2"]>]> {
+                   [Pure, AllTypesMatch<["result", "lhs", "rhs"]>]> {
   let summary = "The `cond ? a : b` ternary operator for vector types";
   let description = [{
     The `cir.vec.ternary` operation represents the C/C++ ternary operator,
@@ -3244,16 +3244,18 @@ def VecTernaryOp : CIR_Op<"vec.ternary",
 
   let arguments = (ins
     CIR_VectorOfIntType:$cond,
-    CIR_VectorType:$vec1,
-		CIR_VectorType:$vec2
+    CIR_VectorType:$lhs,
+    CIR_VectorType:$rhs
   );
 
   let results = (outs CIR_VectorType:$result);
   let assemblyFormat = [{
-    `(` $cond `,` $vec1 `,` $vec2 `)` `:` qualified(type($cond)) `,`
-    qualified(type($vec1)) attr-dict
+    `(` $cond `,` $lhs `,` $rhs `)` `:` qualified(type($cond)) `,`
+    qualified(type($lhs)) attr-dict
   }];
+
   let hasVerifier = 1;
+  let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//

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

@@ -1046,16 +1046,48 @@ LogicalResult cir::VecCreateOp::verify() {
 // VecTernaryOp
 //===----------------------------------------------------------------------===//
 
+OpFoldResult cir::VecTernaryOp::fold(FoldAdaptor adaptor) {
+  mlir::Attribute cond = adaptor.getCond();
+  mlir::Attribute lhs = adaptor.getLhs();
+  mlir::Attribute rhs = adaptor.getRhs();
+
+  if (!mlir::isa_and_nonnull<cir::ConstVectorAttr>(cond) ||
+      !mlir::isa_and_nonnull<cir::ConstVectorAttr>(lhs) ||
+      !mlir::isa_and_nonnull<cir::ConstVectorAttr>(rhs))
+    return {};
+  auto condVec = mlir::cast<cir::ConstVectorAttr>(cond);
+  auto lhsVec = mlir::cast<cir::ConstVectorAttr>(lhs);
+  auto rhsVec = mlir::cast<cir::ConstVectorAttr>(rhs);
+
+  mlir::ArrayAttr condElts = condVec.getElts();
+
+  SmallVector<mlir::Attribute, 16> elements;
+  elements.reserve(condElts.size());
+
+  for (const auto &[idx, condAttr] :
+       llvm::enumerate(condElts.getAsRange<cir::IntAttr>())) {
+    if (condAttr.getSInt()) {
+      elements.push_back(lhsVec.getElts()[idx]);
+    } else {
+      elements.push_back(rhsVec.getElts()[idx]);
+    }
+  }
+
+  cir::VectorType vecTy = getLhs().getType();
+  return cir::ConstVectorAttr::get(
+      vecTy, mlir::ArrayAttr::get(getContext(), elements));
+}
+
 LogicalResult cir::VecTernaryOp::verify() {
   // Verify that the condition operand has the same number of elements as the
   // other operands.  (The automatic verification already checked that all
   // operands are vector types and that the second and third operands are the
   // same type.)
   if (mlir::cast<cir::VectorType>(getCond().getType()).getSize() !=
-      getVec1().getType().getSize()) {
+      getLhs().getType().getSize()) {
     return emitOpError() << ": the number of elements in "
-                         << getCond().getType() << " and "
-                         << getVec1().getType() << " don't match";
+                         << getCond().getType() << " and " << getLhs().getType()
+                         << " don't match";
   }
   return success();
 }

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

clang/lib/CIR/Lowering/DirectToLLVM/LowerToLLVM.cpp

@@ -2008,7 +2008,7 @@ mlir::LogicalResult CIRToLLVMVecTernaryOpLowering::matchAndRewrite(
           op.getCond().getLoc(),
           typeConverter->convertType(op.getCond().getType())));
   rewriter.replaceOpWithNewOp<mlir::LLVM::SelectOp>(
-      op, bitVec, adaptor.getVec1(), adaptor.getVec2());
+      op, bitVec, adaptor.getLhs(), adaptor.getRhs());
   return mlir::success();
 }
 

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

@@ -0,0 +1,19 @@
+// RUN: cir-opt %s -cir-canonicalize -o - | FileCheck %s
+
+!s32i = !cir.int<s, 32>
+
+module {
+  cir.func @vector_ternary_fold_test() -> !cir.vector<!s32i x 4> {
+    %cond = cir.const #cir.const_vector<[#cir.int<1> : !s32i, #cir.int<0> : !s32i, #cir.int<1> : !s32i, #cir.int<0> : !s32i]> : !cir.vector<!s32i x 4>
+    %lhs = 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>
+    %rhs = cir.const #cir.const_vector<[#cir.int<5> : !s32i, #cir.int<6> : !s32i, #cir.int<7> : !s32i, #cir.int<8> : !s32i]> : !cir.vector<!s32i x 4>
+    %res = cir.vec.ternary(%cond, %lhs, %rhs) : !cir.vector<!s32i x 4>, !cir.vector<!s32i x 4>
+    cir.return %res : !cir.vector<!s32i x 4>
+  }
+
+  // [1, 0, 1, 0] ? [1, 2, 3, 4] : [5, 6, 7, 8] Will be fold to [1, 6, 3, 8]
+  // CHECK: cir.func @vector_ternary_fold_test() -> !cir.vector<!s32i x 4> {
+  // CHECK-NEXT: %[[RES:.*]] = cir.const #cir.const_vector<[#cir.int<1> : !s32i, #cir.int<6> : !s32i, #cir.int<3> : !s32i, #cir.int<8> : !s32i]> : !cir.vector<!s32i x 4>
+  // CHECK-NEXT: cir.return %[[RES]] : !cir.vector<!s32i x 4>
+}
+