[Mlir-commits] [mlir] [mlir][spirv] Add folding for SelectOp (PR #85430)

[Jakub Kuderski]

================
@@ -797,6 +797,85 @@ OpFoldResult spirv::LogicalOrOp::fold(FoldAdaptor adaptor) {
   return Attribute();
 }
 
+//===----------------------------------------------------------------------===//
+// spirv.SelectOp
+//===----------------------------------------------------------------------===//
+
+template <class AttrElementT,
+          class ElementValueT = typename AttrElementT::ValueType>
+static Attribute foldSelections(const ElementsAttr &condAttrs,
+                                const ElementsAttr &trueAttrs,
+                                const ElementsAttr &falseAttrs) {
+  auto condsIt = condAttrs.value_begin<BoolAttr>();
+  auto trueAttrsIt = trueAttrs.value_begin<ElementValueT>();
+  auto falseAttrsIt = falseAttrs.value_begin<ElementValueT>();
+
+  SmallVector<ElementValueT, 4> elementResults;
+  elementResults.reserve(condAttrs.getNumElements());
+  for (size_t i = 0, e = condAttrs.getNumElements(); i < e;
+       ++i, ++condsIt, ++trueAttrsIt, ++falseAttrsIt) {
+    if ((*condsIt).getValue()) // If Condition then take Object 1
+      elementResults.push_back(*trueAttrsIt);
+    else // Else take Object 2
+      elementResults.push_back(*falseAttrsIt);
+  }
+
+  auto resultType = trueAttrs.getType();
+  return DenseElementsAttr::get(cast<ShapedType>(resultType), elementResults);
+}
+
+static Attribute foldSelectOp(llvm::ArrayRef<Attribute> operands) {
+  auto condAttrs = llvm::dyn_cast<ElementsAttr>(operands[0]);
+  auto trueAttrs = llvm::dyn_cast<ElementsAttr>(operands[1]);
+  auto falseAttrs = llvm::dyn_cast<ElementsAttr>(operands[2]);
----------------
kuhar wrote:

nit
```suggestion
  auto condAttrs = dyn_cast<ElementsAttr>(operands[0]);
  auto trueAttrs = dyn_cast<ElementsAttr>(operands[1]);
  auto falseAttrs = dyn_cast<ElementsAttr>(operands[2]);
```

https://github.com/llvm/llvm-project/pull/85430