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

[Jakub Kuderski]

================
@@ -801,16 +801,18 @@ OpFoldResult spirv::LogicalOrOp::fold(FoldAdaptor adaptor) {
 // spirv.SelectOp
 //===----------------------------------------------------------------------===//
 
-template <class AttrElementT,
-          class ElementValueT = typename AttrElementT::ValueType>
-static Attribute foldSelections(const ElementsAttr &condAttrs,
-                                const ElementsAttr &trueAttrs,
-                                const ElementsAttr &falseAttrs) {
+static Attribute foldSelectOp(ArrayRef<Attribute> operands) {
+  auto condAttrs = dyn_cast<DenseElementsAttr>(operands[0]);
+  auto trueAttrs = dyn_cast<DenseElementsAttr>(operands[1]);
+  auto falseAttrs = dyn_cast<DenseElementsAttr>(operands[2]);
+  if (!condAttrs || !trueAttrs || !falseAttrs)
+    return Attribute();
+
   auto condsIt = condAttrs.value_begin<BoolAttr>();
-  auto trueAttrsIt = trueAttrs.value_begin<ElementValueT>();
-  auto falseAttrsIt = falseAttrs.value_begin<ElementValueT>();
+  auto trueAttrsIt = trueAttrs.value_begin<Attribute>();
+  auto falseAttrsIt = falseAttrs.value_begin<Attribute>();
 
-  SmallVector<ElementValueT, 4> elementResults;
+  SmallVector<Attribute, 4> elementResults;
   elementResults.reserve(condAttrs.getNumElements());
   for (size_t i = 0, e = condAttrs.getNumElements(); i < e;
----------------
kuhar wrote:

nit
```suggestion
  size_t numElements = condAttrs.getNumElements();
  elementResults.reserve(numElements);
  for (size_t i = 0; i != numElements;
```

or if this works (iterators are sometimes weird and don't compose with all the utility functions):
```c++
auto elementResults = llvm::to_vector<4>(trueAttrs.getValues<Attribute>());
for (auto [result, cond, falseRes] : llvm::zip_equal(elementResults, condAttrs.getValues<BoolAttribute>(), falseAttr.getValues<Attribute>())) {
  if (cond.getValue())
    result = falseRes;
}
```

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