[CIR] Implement folder for VecTernaryOp #142946

AmrDeveloper · 2025-06-05T11:11:36Z

This change adds a folder for the VecTernaryOp

Issue #136487

llvmbot · 2025-06-05T11:12:10Z

@llvm/pr-subscribers-clang

Author: Amr Hesham (AmrDeveloper)

Changes

This change adds a folder for the VecTernaryOp

Issue #136487

Full diff: https://github.com/llvm/llvm-project/pull/142946.diff

4 Files Affected:

(modified) clang/include/clang/CIR/Dialect/IR/CIROps.td (+2)
(modified) clang/lib/CIR/Dialect/IR/CIRDialect.cpp (+35)
(modified) clang/lib/CIR/Dialect/Transforms/CIRCanonicalize.cpp (+3-3)
(added) clang/test/CIR/Transforms/vector-ternary-fold.cir (+20)

diff --git a/clang/include/clang/CIR/Dialect/IR/CIROps.td b/clang/include/clang/CIR/Dialect/IR/CIROps.td
index 00878f7dd8ed7..eb439f7aa1527 100644
--- a/clang/include/clang/CIR/Dialect/IR/CIROps.td
+++ b/clang/include/clang/CIR/Dialect/IR/CIROps.td
@@ -2228,7 +2228,9 @@ def VecTernaryOp : CIR_Op<"vec.ternary",
     `(` $cond `,` $lhs`,` $rhs `)` `:` qualified(type($cond)) `,`
     qualified(type($lhs)) attr-dict
   }];
+
   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 fa7fb592a3cd6..f585254d3340b 100644
--- a/clang/lib/CIR/Dialect/IR/CIRDialect.cpp
+++ b/clang/lib/CIR/Dialect/IR/CIRDialect.cpp
@@ -1638,6 +1638,41 @@ LogicalResult cir::VecTernaryOp::verify() {
   return success();
 }
 
+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)) {
+    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]);
+        continue;
+      }
+
+      elements.push_back(rhsVec.getElts()[idx]);
+    }
+
+    cir::VectorType vecTy = getLhs().getType();
+    return cir::ConstVectorAttr::get(
+        vecTy, mlir::ArrayAttr::get(getContext(), elements));
+  }
+
+  return {};
+}
+
 //===----------------------------------------------------------------------===//
 // TableGen'd op method definitions
 //===----------------------------------------------------------------------===//
diff --git a/clang/lib/CIR/Dialect/Transforms/CIRCanonicalize.cpp b/clang/lib/CIR/Dialect/Transforms/CIRCanonicalize.cpp
index 7d03e374c27e8..aa3e97033cdda 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, VecExtractOp and VecShuffleDynamicOp are here to perform
-    // a manual `fold` in applyOpPatternsGreedily.
+    // CastOp, UnaryOp, VecExtractOp, VecShuffleDynamicOp and VecTernaryOp are
+    // here to perform a manual `fold` in applyOpPatternsGreedily.
     if (isa<BrOp, BrCondOp, CastOp, ScopeOp, SwitchOp, SelectOp, UnaryOp,
-            VecExtractOp, VecShuffleDynamicOp>(op))
+            VecExtractOp, VecShuffleDynamicOp, VecTernaryOp>(op))
       ops.push_back(op);
   });
 
diff --git a/clang/test/CIR/Transforms/vector-ternary-fold.cir b/clang/test/CIR/Transforms/vector-ternary-fold.cir
new file mode 100644
index 0000000000000..f2e18576da74b
--- /dev/null
+++ b/clang/test/CIR/Transforms/vector-ternary-fold.cir
@@ -0,0 +1,20 @@
+// RUN: cir-opt %s -cir-canonicalize -o - | FileCheck %s
+
+!s32i = !cir.int<s, 32>
+
+module {
+  cir.func @vector_ternary_fold_test() -> !cir.vector<4 x !s32i> {
+    %cond = cir.const #cir.const_vector<[#cir.int<1> : !s32i, #cir.int<0> : !s32i, #cir.int<1> : !s32i, #cir.int<0> : !s32i]> : !cir.vector<4 x !s32i>
+    %lhs = 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>
+    %rhs = cir.const #cir.const_vector<[#cir.int<5> : !s32i, #cir.int<6> : !s32i, #cir.int<7> : !s32i, #cir.int<8> : !s32i]> : !cir.vector<4 x !s32i>
+    %res = cir.vec.ternary(%cond, %lhs, %rhs) : !cir.vector<4 x !s32i>, !cir.vector<4 x !s32i>
+    cir.return %res : !cir.vector<4 x !s32i>
+  }
+
+  // [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<4 x !s32i> {
+  // 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<4 x !s32i>
+  // CHECK-NEXT: cir.return %[[RES]] : !cir.vector<4 x !s32i>
+}
+
+

llvmbot · 2025-06-05T11:12:10Z

@llvm/pr-subscribers-clangir

Author: Amr Hesham (AmrDeveloper)

Changes

This change adds a folder for the VecTernaryOp

Issue #136487

Full diff: https://github.com/llvm/llvm-project/pull/142946.diff

4 Files Affected:

(modified) clang/include/clang/CIR/Dialect/IR/CIROps.td (+2)
(modified) clang/lib/CIR/Dialect/IR/CIRDialect.cpp (+35)
(modified) clang/lib/CIR/Dialect/Transforms/CIRCanonicalize.cpp (+3-3)
(added) clang/test/CIR/Transforms/vector-ternary-fold.cir (+20)

diff --git a/clang/include/clang/CIR/Dialect/IR/CIROps.td b/clang/include/clang/CIR/Dialect/IR/CIROps.td
index 00878f7dd8ed7..eb439f7aa1527 100644
--- a/clang/include/clang/CIR/Dialect/IR/CIROps.td
+++ b/clang/include/clang/CIR/Dialect/IR/CIROps.td
@@ -2228,7 +2228,9 @@ def VecTernaryOp : CIR_Op<"vec.ternary",
     `(` $cond `,` $lhs`,` $rhs `)` `:` qualified(type($cond)) `,`
     qualified(type($lhs)) attr-dict
   }];
+
   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 fa7fb592a3cd6..f585254d3340b 100644
--- a/clang/lib/CIR/Dialect/IR/CIRDialect.cpp
+++ b/clang/lib/CIR/Dialect/IR/CIRDialect.cpp
@@ -1638,6 +1638,41 @@ LogicalResult cir::VecTernaryOp::verify() {
   return success();
 }
 
+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)) {
+    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]);
+        continue;
+      }
+
+      elements.push_back(rhsVec.getElts()[idx]);
+    }
+
+    cir::VectorType vecTy = getLhs().getType();
+    return cir::ConstVectorAttr::get(
+        vecTy, mlir::ArrayAttr::get(getContext(), elements));
+  }
+
+  return {};
+}
+
 //===----------------------------------------------------------------------===//
 // TableGen'd op method definitions
 //===----------------------------------------------------------------------===//
diff --git a/clang/lib/CIR/Dialect/Transforms/CIRCanonicalize.cpp b/clang/lib/CIR/Dialect/Transforms/CIRCanonicalize.cpp
index 7d03e374c27e8..aa3e97033cdda 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, VecExtractOp and VecShuffleDynamicOp are here to perform
-    // a manual `fold` in applyOpPatternsGreedily.
+    // CastOp, UnaryOp, VecExtractOp, VecShuffleDynamicOp and VecTernaryOp are
+    // here to perform a manual `fold` in applyOpPatternsGreedily.
     if (isa<BrOp, BrCondOp, CastOp, ScopeOp, SwitchOp, SelectOp, UnaryOp,
-            VecExtractOp, VecShuffleDynamicOp>(op))
+            VecExtractOp, VecShuffleDynamicOp, VecTernaryOp>(op))
       ops.push_back(op);
   });
 
diff --git a/clang/test/CIR/Transforms/vector-ternary-fold.cir b/clang/test/CIR/Transforms/vector-ternary-fold.cir
new file mode 100644
index 0000000000000..f2e18576da74b
--- /dev/null
+++ b/clang/test/CIR/Transforms/vector-ternary-fold.cir
@@ -0,0 +1,20 @@
+// RUN: cir-opt %s -cir-canonicalize -o - | FileCheck %s
+
+!s32i = !cir.int<s, 32>
+
+module {
+  cir.func @vector_ternary_fold_test() -> !cir.vector<4 x !s32i> {
+    %cond = cir.const #cir.const_vector<[#cir.int<1> : !s32i, #cir.int<0> : !s32i, #cir.int<1> : !s32i, #cir.int<0> : !s32i]> : !cir.vector<4 x !s32i>
+    %lhs = 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>
+    %rhs = cir.const #cir.const_vector<[#cir.int<5> : !s32i, #cir.int<6> : !s32i, #cir.int<7> : !s32i, #cir.int<8> : !s32i]> : !cir.vector<4 x !s32i>
+    %res = cir.vec.ternary(%cond, %lhs, %rhs) : !cir.vector<4 x !s32i>, !cir.vector<4 x !s32i>
+    cir.return %res : !cir.vector<4 x !s32i>
+  }
+
+  // [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<4 x !s32i> {
+  // 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<4 x !s32i>
+  // CHECK-NEXT: cir.return %[[RES]] : !cir.vector<4 x !s32i>
+}
+
+

xlauko

lgtm

andykaylor

This looks good. I just have a few minor suggestions.

andykaylor · 2025-06-05T20:50:56Z

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

+         llvm::enumerate(condElts.getAsRange<cir::IntAttr>())) {
+      if (condAttr.getSInt()) {
+        elements.push_back(lhsVec.getElts()[idx]);
+        continue;


I think this would be more natural as an if-else rather than using continue.

andykaylor · 2025-06-05T20:57:24Z

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

+  if (mlir::isa_and_nonnull<cir::ConstVectorAttr>(cond) &&
+      mlir::isa_and_nonnull<cir::ConstVectorAttr>(lhs) &&
+      mlir::isa_and_nonnull<cir::ConstVectorAttr>(rhs)) {


Suggested change

if (mlir::isa_and_nonnull<cir::ConstVectorAttr>(cond) &&

mlir::isa_and_nonnull<cir::ConstVectorAttr>(lhs) &&

mlir::isa_and_nonnull<cir::ConstVectorAttr>(rhs)) {

if (!mlir::isa_and_nonnull<cir::ConstVectorAttr>(cond) ||

!mlir::isa_and_nonnull<cir::ConstVectorAttr>(lhs) ||

!mlir::isa_and_nonnull<cir::ConstVectorAttr>(rhs))

return {};

andykaylor · 2025-06-05T21:03:39Z

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

-    // CastOp, UnaryOp, VecExtractOp and VecShuffleDynamicOp are here to perform
-    // a manual `fold` in applyOpPatternsGreedily.
+    // CastOp, UnaryOp, VecExtractOp, VecShuffleDynamicOp and VecTernaryOp are
+    // here to perform a manual `fold` in applyOpPatternsGreedily.


Amongst the operations that are here to perform a manual fold are.... (This is Monty Python joke, but we really should just make this a general comment that many of the operations in this list are just here so we can fold them.)

I think we can split them into two if statements, one for the ops that are here to perform manual fold and another for the other ops

if (isa<BrOp, BrCondOp, ScopeOp, SwitchOp, SelectOp>(op)) ops.push_back(op); // Operations to perform manual `fold` in applyOpPatternsGreedily. if (isa<CastOp, UnaryOp, VecExtractOp, VecShuffleDynamicOp, VecTernaryOp>( op)) ops.push_back(op);

That's not really necessary. I think we've passed the point where anyone is going to be checking each one to see why it's here. Also, some operations may be handled explicitly and also have a folder.

Yes, you are right, i will convert it to general comment

This change adds a folder for the VecTernaryOp Issue llvm#136487

[CIR] Implement folder for VecTernaryOp

ac8277b

AmrDeveloper requested review from xlauko, mmha, erichkeane and andykaylor June 5, 2025 11:11

AmrDeveloper requested review from lanza and bcardosolopes as code owners June 5, 2025 11:11

llvmbot added clang Clang issues not falling into any other category ClangIR Anything related to the ClangIR project labels Jun 5, 2025

xlauko approved these changes Jun 5, 2025

View reviewed changes

andykaylor approved these changes Jun 5, 2025

View reviewed changes

AmrDeveloper added 2 commits June 5, 2025 23:28

Address code review comments

9263706

Make the comment more general for ops that here for fold

fc8b21c

AmrDeveloper merged commit 995d74f into llvm:main Jun 7, 2025
7 checks passed

tomtor pushed a commit to tomtor/llvm-project that referenced this pull request Jun 14, 2025

[CIR] Implement folder for VecTernaryOp (llvm#142946)

0cdbf7f

This change adds a folder for the VecTernaryOp Issue llvm#136487

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[CIR] Implement folder for VecTernaryOp #142946

[CIR] Implement folder for VecTernaryOp #142946

Uh oh!

AmrDeveloper commented Jun 5, 2025

Uh oh!

llvmbot commented Jun 5, 2025

Uh oh!

llvmbot commented Jun 5, 2025

Uh oh!

xlauko left a comment

Uh oh!

andykaylor left a comment

Uh oh!

andykaylor Jun 5, 2025

Uh oh!

andykaylor Jun 5, 2025

Uh oh!

andykaylor Jun 5, 2025

Uh oh!

AmrDeveloper Jun 5, 2025 •

edited

Loading

Uh oh!

andykaylor Jun 5, 2025

Uh oh!

AmrDeveloper Jun 5, 2025

Uh oh!

Uh oh!

Uh oh!

[CIR] Implement folder for VecTernaryOp #142946

[CIR] Implement folder for VecTernaryOp #142946

Uh oh!

Conversation

AmrDeveloper commented Jun 5, 2025

Uh oh!

llvmbot commented Jun 5, 2025

Uh oh!

llvmbot commented Jun 5, 2025

Uh oh!

xlauko left a comment

Choose a reason for hiding this comment

Uh oh!

andykaylor left a comment

Choose a reason for hiding this comment

Uh oh!

andykaylor Jun 5, 2025

Choose a reason for hiding this comment

Uh oh!

andykaylor Jun 5, 2025

Choose a reason for hiding this comment

Uh oh!

andykaylor Jun 5, 2025

Choose a reason for hiding this comment

Uh oh!

AmrDeveloper Jun 5, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

andykaylor Jun 5, 2025

Choose a reason for hiding this comment

Uh oh!

AmrDeveloper Jun 5, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

AmrDeveloper Jun 5, 2025 •

edited

Loading